llvmtypes.hs

module Main where

import Control.Applicative
import Control.Exception
import Data.List
import Data.Char
import qualified Data.Map as Map
import System.IO
import System.IO.Error
import System.Environment
import System.Process

data Typ =
    V | Opaque | I1 | I8 | I16 | I32 | I64 | F32 | F64 | Ref String |
    Arr { arrCount :: Word, arrElemTyp :: Typ } |
    Stru { struMembers :: [Typ] } | PStru { struMembers :: [Typ] } |
    Ptr Typ | Fun { retTyp :: Typ, argTyp :: [Typ], isVarArg :: Bool }
  deriving (Read, Show, Eq)

data SzAl = SzAl { saSize :: Word, saAlign :: Word }
  deriving (Read, Show, Eq, Ord)

isIdChar :: Char -> Bool
isIdChar = (`elem` "._"++['0'..'9']++['a'..'z']++['A'..'Z'])

stE :: String -> String
stE = dropWhile isSpace

literal :: String -> String -> [((), String)]
literal p i =
  case stripPrefix p i of
    Just x  -> return ((), stE x)
    Nothing -> fail "literal match fail"

ident :: String -> [(String, String)]
ident (x:xs)
    | isIdChar x =
        let (v, r) = span isIdChar xs
        in return ((x:v), stE r)
    | otherwise  = fail "ident is empty"
ident [] = fail "unexpected EOF for ident"

typeAtom :: String -> [(Typ, String)]
typeAtom ('o':'p':'a':'q':'u':'e':xs) = return (Opaque, stE xs)
typeAtom ('v':'o':'i':'d':xs) = return (V, stE xs)
typeAtom ('i':'8':xs) = return (I8, stE xs)
typeAtom ('i':'1':'6':xs) = return (I16, stE xs)
typeAtom ('i':'3':'2':xs) = return (I32, stE xs)
typeAtom ('i':'6':'4':xs) = return (I64, stE xs)
typeAtom ('f':'l':'o':'a':'t':xs) = return (F32, stE xs)
typeAtom ('d':'o':'u':'b':'l':'e':xs) = return (F64, stE xs)
typeAtom ('i':'1':xs) = return (I1, stE xs)
typeAtom ('%':xs) = do
  (n, r) <- ident (stE xs)
  return (Ref n, r)
typeAtom ('(':xs) = do
  (v, r) <- typeExpr (stE xs)
  (_, r) <- literal ")" r
  return (v, r)
typeAtom ('[':xs) = do
  (c, r) <- reads (stE xs)
  (_, r) <- literal "x" (stE r)
  (v, r) <- typeExpr r
  (_, r) <- literal "]" r
  return (Arr c v, r)
typeAtom ('<':xs) = do
  (_, r) <- literal "{" (stE xs)
  let m0 ('}':xs) = mf (stE xs) id
      m0 r = m r id
      mf r vs = do
        (_, r) <- literal ">" r
        let vf = vs []
        vf `seq` return (PStru vf, r)
      m r vs = do
        (v, r) <- typeExpr r
        let vs' = vs . (v:)
        case r of
          (',':xs) -> m (stE xs) vs'
          ('}':xs) -> mf (stE xs) vs'
          _ -> fail "malformed packed product type"
  m0 r
typeAtom ('{':xs) = m0 (stE xs) where
  m0 ('}':xs) = return (Stru [], stE xs)
  m0 r = m r id
  m r vs = do
    (v, r) <- typeExpr r
    let vs' = vs . (v:)
    case r of
      (',':xs) -> m (stE xs) vs'
      ('}':xs) ->
        let vf = vs' []
        in vf `seq` return (Stru vf, stE xs)
      _ -> fail "malformed product type"
typeAtom _ = fail "not a type"

typeExpr :: String -> [(Typ, String)]
typeExpr inp = do
  (v, r) <- typeAtom inp
  let m k ('*':xs) = m (Ptr k) (stE xs)
      m k ('(':xs) = w (stE xs) id where
        fin r vs b = 
          let vf = vs []
          in vf `seq` m (Fun k vf b) (stE r)
        fin0 r vs = fin r vs False
        finV r vs = do
          (_, r) <- literal ")" r
          fin r vs True
        nxt ('.':'.':'.':xs) vs = finV (stE xs) vs
        nxt r vs = do
          (v, r) <- typeExpr r
          let vs' = vs . (v:)
          case r of
            (',':xs) -> nxt (stE xs) vs'
            (')':xs) -> fin0 (stE xs) vs'
            _       -> fail "expected ',' or ')'"
        w ('.':'.':'.':xs) vs = finV (stE xs) vs
        w (')':xs) vs = fin0 (stE xs) vs
        w      r   vs = nxt r vs
      m k r = return (k, r)
  m v r

unref :: Applicative f => (String -> f Typ) -> Typ -> f Typ
unref f (Ref x) = f x
unref f (Arr c e) = Arr c <$> unref f e
unref f (Stru x) = Stru <$> traverse (unref f) x
unref f (PStru x) = PStru <$> traverse (unref f) x
unref f (Ptr t) = Ptr <$> unref f t
unref f (Fun rt ats b) = liftA2 (\x y -> Fun x y b) (unref f rt) (traverse (unref f) ats)
unref _ t = pure t

typeDecl :: String -> [(String, Typ, String)]
typeDecl ('%':xs) = do
  (n, r) <- ident xs
  (_, r) <- literal "=" r
  (_, r) <- literal "type" r
  (t, r) <- typeExpr r
  return (n, t, r)
typeDecl _ = fail "not a type declaration"

varDecl :: String -> [(String, Typ, String)]
varDecl ('@':xs) = do
  (n, r) <- ident xs
  (_, r) <- literal "=" r
  let m r@(_:_) = m0 r <|> m (stK r)
      m _ = fail "unexpected EOF"
      stK = stE . dropWhile (not . isSpace)
      m0 = fmap (\(t, r) -> (n, t, r)) . typeExpr
  m (stK r)
varDecl _ = fail "not a var declaration"

parseType :: String -> Either String Typ
parseType inp =
  case [ x | (x, "") <- typeExpr (stE inp) ] of
    [x] -> Right x
    []  -> Left "parseType: no parse"
    _   -> Left "parseType: ambiguous parse"

sizeOf :: SzAl -> Typ -> SzAl
sizeOf _ V = SzAl 1 1
sizeOf _ I1 = SzAl 1 1
sizeOf _ I8 = SzAl 1 1
sizeOf _ I16 = SzAl 2 2
sizeOf _ I32 = SzAl 4 4
sizeOf _ I64 = SzAl 8 8
sizeOf _ F32 = SzAl 4 4
sizeOf _ F64 = SzAl 8 8
sizeOf p (Arr c e) = SzAl (c * s) a where SzAl s a = sizeOf p e
sizeOf p (Stru m) =
  let sa = map (sizeOf p) m
      f (SzAl s0 a0) (SzAl s a) =
        let k = (s0 + a - 1) `quot` a
        in SzAl (a * k + s) (lcm a0 a)
      SzAl s' a' = foldl' f (SzAl 0 1) sa
  in f (SzAl s' 1) (SzAl 0 a')
sizeOf p (PStru m) = SzAl s 1 where
  s = foldr ((+) . saSize . sizeOf p) 0 m
sizeOf p (Ptr _) = p
sizeOf _ _ = SzAl 0 1

sizeOf32 :: Typ -> SzAl
sizeOf32 = sizeOf (SzAl 4 4)

sizeOf64 :: Typ -> SzAl
sizeOf64 = sizeOf (SzAl 8 8)

type TypMap = Map.Map String Typ
data Action = NewVar String Typ | UpdateReg String Typ | Failure String | Nop
  deriving (Read, Show, Eq)

procLine :: String -> Action
procLine r =
  case r of
    '%':_ ->
      case typeDecl r of
        (n, t, _):_ -> UpdateReg n t
        _           -> Failure "unknown type decl"
    '@':_ ->
      case varDecl r of
        (n, t, _):_ -> NewVar n t
        _           -> Failure "unknown var decl"
    _ -> Nop

procLLFile :: String -> Handle -> IO ()
procLLFile fn fh = procIter (1 :: Word) Map.empty `catch` handleIt where
  sizeOf' = saSize . sizeOf32
  fail' = Left
  resolve tm hist n
    | n `elem` hist =
        fail' $ "Cycle! " ++ show hist
    | otherwise =
        case Map.lookup n tm of
          Just (Ref x) -> resolve tm (n:hist) x
          Just t -> return t
          _ -> fail' $ "Unknown type " ++ n
  handleIt :: IOError -> IO ()
  handleIt e
    | isEOFError e = return ()
    | otherwise    = throwIO e
  rptRoot = hPutStr stderr . ((fn ++ ": ") ++)
  procIter lno tm = do
    line <- hGetLine fh
    let report str =
          rptRoot $ shows lno $ ": " ++ str ++ ": " ++ line ++ "\n"
    case procLine line of
      NewVar _n t -> do
        case unref (resolve tm []) t of
          Left e -> report e
          Right t' -> putStrLn v where
            v = shows (sizeOf' t') (' ':fn ++ (':':line))
        procIter (lno+1) tm
      UpdateReg n t -> procIter (lno+1) (Map.insert n t tm)
      Failure e -> report e *> procIter (lno+1) tm
      Nop -> procIter (lno+1) tm

procFile :: String -> Handle -> IO ()
procFile fn fh =
  withCreateProcess (proc "llvm-dis" ["-o=-"]) {
    std_in = UseHandle fh,
    std_out = CreatePipe
  } $ \_ ofhm _ _ -> do
    let Just ofh = ofhm
    procLLFile fn ofh

main :: IO ()
main = getArgs >>= \args ->
  case args of
   [] -> procFile "-" stdin
   _ -> mapM_ (\fn -> withFile fn ReadMode (procFile fn)) args