module Compile (compile) where

import Control.Monad.State hiding (State)
import Data.List
import Language.C
import Language.C.Data.Ident

import Error
import Model hiding (CInteger)
import qualified Model as M
import Utils

-- | Compiles a program to a model for analysis.
compile :: String -> [CTranslUnit] -> IO Model
compile function units = do
  m <- execStateT (evalStat initEnv $ rewrite function units) initMDB
  return (model m) { actions = reverse $ actions (model m) }

none :: NodeInfo
none = internalNode

-- | Rewrites a program to a single statement.  Requires no recursive functions and unique declarations for all top level declarations.
rewrite :: String -> [CTranslUnit] -> CStat
rewrite name units = if not $ null asms
  then notSupported (head asms) "inline assembly"
  else if not $ null duplicateNames
    then error $ "duplicate top-level names found (hiding names with static is not supported): " ++ show duplicateNames
    else CCompound [] (vars ++ funcs ++ [CBlockStmt call]) none
  where
  items = [ a | CTranslUnit items _ <- units, a <- items ]
  varDefs  = [ a | CDeclExt a <- items ]
  vars  = [ CBlockDecl (CDecl (CStorageSpec (CStatic none) : specs) a b) | CDecl specs a b <- varDefs ]  -- Make top level vars static.
  funcDefs = [ a | CFDefExt a <- items ]
  funcs = map CNestedFunDef $ sortFunctions funcDefs
  asms  = [ a | CAsmExt  a <- items ]
  call :: CStat
  call = CExpr (Just $ CCall (CVar (Ident name 0 none) none) [] none) none
  duplicateNames = duplicates $ [ name | f <- funcDefs, let (_, (Ident name _ _), _, _) = functionInfo f ] ++ concat [ [ name | (Just (CDeclr (Just (Ident name _ _)) _ _ _ _), _, _) <- a ] | CDecl _ a _ <- varDefs ]

duplicates :: Eq a => [a] -> [a]
duplicates [] = []
duplicates (a:b) | elem a b  = a : duplicates b
                 | otherwise =     duplicates b

type M = StateT MDB IO

data MDB = MDB
  { nextId  :: Int
  , stack   :: [Ident]
  , enabled :: E
  , model   :: Model
  }

initMDB :: MDB
initMDB = MDB
  { nextId  = 0
  , stack   = []
  , enabled = true
  , model   = Model { variables = [], actions = [] }
  }



-- | Environment for resolving identifiers.
data Env = Env
  { values  :: [Value]
  }

data Value
  = EnvFunction String Function
  | EnvVariable String V

data Function = Function Int ([E] -> M E)

-- | Looks of a variable in the environment.
variable :: Env -> Ident -> V
variable env i@(Ident name _ _) = if null m then err i $ "variable \"" ++ name ++ "\" not found" else head m
  where
  m = [ a | EnvVariable n a <- values env, n == name ]

-- | Looks of a function in the environment.
function :: Env -> Ident -> Function
function env i@(Ident name _ _) = if null m then err i $ "function \"" ++ name ++ "\" not found" else head m
  where
  m = [ a | EnvFunction n a <- values env, n == name ]

-- | Creates a branch.
branch :: Position -> E -> M a -> M a -> M (a, a)
branch n a onTrue onFalse = do
  m1 <- get
  put m1 { enabled = And (enabled m1) a n }
  r1 <- onTrue
  m2 <- get
  put m2 { enabled = And (enabled m1) (Not a n) n }
  r2 <- onFalse
  m3 <- get
  put m3 { enabled = enabled m1 }
  return (r1, r2)

-- | Push an identifier onto the call stack, do something, then pop it off.
callStack :: Ident -> M a -> M a
callStack id a = do
  m <- get
  put m { stack = id : stack m }
  a <- a
  m <- get
  put m { stack = tail $ stack m }
  return a

callPath :: M ([String], Position)
callPath = do
  m <- get
  let s = stack m
  return ([ n | Ident n _ _ <- reverse s ], posOf $ head s)

-- | The initial environment defines the assert and assume functions.
initEnv :: Env
initEnv = Env
  { values =
    [ EnvFunction "assert" $ Function 1 assert
    , EnvFunction "assume" $ Function 1 assume
    ]
  }
  where
  assert a = do
    (s, n) <- callPath
    m <- get
    let x = imply (enabled m) (head a) n
    newAction $ Assert x s n
    return x
  assume a = do
    (s, n) <- callPath
    m <- get
    let x = imply (enabled m) (head a) n
    newAction $ Assume x s n
    return x

-- | Adds new action.
newAction :: Action -> M ()
newAction a = do
  m <- get
  put m { model = (model m) { actions = a : actions (model m) }}

-- | Adds new variable.
addVar :: Env -> V -> M Env
addVar env a = do
  case a of
    State a -> do
      m <- get
      put m { model = (model m) { variables = if elem a (variables $ model m) then variables $ model m else a : variables (model m) }}
    _ -> return ()
  return env { values = EnvVariable name a : values env }
  where
  name = case a of
    State (VS a _ _ _)  -> a
    Volatile a _ _   -> a
    Local    a _ _ _ -> a
    Tmp      _ _ _   -> error "Compile.addVar: should not call addVar with Tmp"



evalStat :: Env -> CStat -> M ()
evalStat env a = case a of
  CLabel i a [] _ -> callStack i $ evalStat env a
  CExpr Nothing _ -> return ()
  CExpr (Just a) _ -> evalExpr env a >> return ()
  CCompound ids items _ -> f ids
    where
    f :: [Ident] -> M ()
    f [] = foldM evalBlockItem env items >> return ()
    f (a:b) = callStack a $ f b
  CIf a b Nothing n -> evalStat env $ CIf a b (Just $ CCompound [] [] n) n
  CIf a b (Just c) n -> do
    a <- evalExpr env a
    branch (posOf n) a (evalStat env b) (evalStat env c)
    return ()
  _ -> notSupported a "statement"

evalBlockItem :: Env -> CBlockItem -> M Env
evalBlockItem env a = case a of
  CBlockStmt a    -> evalStat env a >> return env
  CBlockDecl a    -> evalDecl env a
  CNestedFunDef a -> evalFunc env a

evalExpr :: Env -> CExpr -> M E
evalExpr env a = latch (posOf a) =<< case a of
  CAssign op a b n -> case op of
    CAssignOp -> do
      a' <- evalExpr env a
      b' <- evalExpr env b
      case a' of
        Var v -> assign False (posOf n) v b' >> return a'
        _ -> unexpected a "non variable in left hand of assignment"
    CMulAssOp -> f CMulOp
    CDivAssOp -> f CDivOp
    CRmdAssOp -> f CRmdOp
    CAddAssOp -> f CAddOp
    CSubAssOp -> f CSubOp
    CShlAssOp -> f CShlOp
    CShrAssOp -> f CShrOp
    CAndAssOp -> f CAndOp
    CXorAssOp -> f CXorOp
    COrAssOp  -> f COrOp
    where
    f :: CBinaryOp -> M E
    f op = evalExpr env (CAssign CAssignOp a (CBinary op a b n) n)
  CCond a (Just b) c n -> do
    a <- evalExpr env a
    (b, c) <- branch (posOf n) a (evalExpr env b) (evalExpr env c)
    return $ Mux a b c $ posOf n
  CCond a Nothing b n -> do
    a <- evalExpr env a
    (a, b) <- branch (posOf n) a (return a) (evalExpr env b)
    return $ Mux a a b $ posOf n
  CBinary op a b n -> case op of
    CMulOp -> f Mul
    CDivOp -> f Div
    CRmdOp -> f Mod
    CAddOp -> f Add
    CSubOp -> f Sub
    CShlOp -> notSupported a "(<<)"
    CShrOp -> notSupported a "(>>)"
    CLeOp  -> f Lt
    CGrOp  -> f $ \ a b n -> Lt b a n
    CLeqOp -> f $ \ a b n -> Not (Lt b a n) n
    CGeqOp -> f $ \ a b n -> Not (Lt a b n) n
    CEqOp  -> f Eq
    CNeqOp -> f $ \ a b n -> Not (Eq a b n) n
    CAndOp -> notSupported a "(&)"
    CXorOp -> notSupported a "(^)"
    COrOp  -> notSupported a "(|)"
    CLndOp -> do
      a <- evalExpr env a
      (b, a) <- branch n' a (evalExpr env b) (return a)
      return $ Mux a b a n'
    CLorOp -> do
      a <- evalExpr env a
      (a, b) <- branch n' a (return a) (evalExpr env b)
      return $ Mux a a b n'
    where
    n' = posOf n
    f :: (E -> E -> Position -> E) -> M E
    f op = do
      a <- evalExpr env a
      b <- evalExpr env b
      return $ op a b n'

  CUnary op a n -> do
    a <- evalExpr env a
    case (op, a) of
      (CPreIncOp, Var a) -> do
        assign False p a $ Add (Var a) one p
        return $ Var a
      (CPreDecOp, Var a) -> do
        assign False p a $ Sub (Var a) one p
        return $ Var a
      (CPostIncOp, Var a) -> do
        b <- latch p $ Var a
        assign False p a $ Add (Var a) one p
        return b
      (CPostDecOp, Var a) -> do
        b <- latch p $ Var a
        assign False p a $ Sub (Var a) one p
        return b
      (CPlusOp, a) -> return a
      (CMinOp,  a) -> return $ Sub zero a p
      (CNegOp,  a) -> return $ Not a p
      _ -> notSupported n "unary operator"
    where
    p = posOf n
    one  = Const $ M.CInteger 1 p
    zero = Const $ M.CInteger 0 p

  CCall (CVar f _) args _ -> do
    args <- mapM (evalExpr env) args
    when (arity /= length args) $ unexpected f $ "function called with " ++ show (length args) ++ " arguments, but defined with " ++ show arity ++ " arguments"
    callStack f $ func args
    where
    Function arity func = function env f

  CCall _ _ _ -> notSupported a "non named function references"

  CVar i _ -> return $ Var $ variable env i
  CConst a -> case a of
    CIntConst (CInteger a _ _) n -> return $ Const $ M.CInteger a $ posOf n
    CFloatConst (CFloat a) n -> return $ Const $ CRational (toRational (read a :: Double)) $ posOf n
    _ -> notSupported a "char or string constant"
  _ -> notSupported a "expression"
    
evalDecl :: Env -> CDecl -> M Env
evalDecl env d@(CDecl specs decls _) = if isExtern typInfo then return env else foldM evalDecl' env decls
  where
  (typInfo, typ) = typeInfo specs
  evalDecl' :: Env -> (Maybe CDeclr, Maybe CInit, Maybe CExpr) -> M Env
  evalDecl' env (a, b, c) = case a of
    Just (CDeclr (Just i@(Ident name _ n)) [] Nothing [] _) -> case (b, c) of
      (Nothing, Nothing) -> evalDecl' env (a, Just $ CInitExpr (CConst (CIntConst (cInteger 0) n)) n, Nothing)
      (Just (CInitExpr (CConst const) n'), Nothing) | isStatic typInfo && not (isVolatile typInfo) -> addVar env v
        where
        v = State $ VS name typ init $ posOf n
        init = case typ of
          Void -> unexpected d "void type for variable declaration"
          Bool -> CBool (cInt /= 0) $ posOf n'
          Integer _  -> M.CInteger cInt $ posOf n'
          Rational _ -> CRational cRat $ posOf n'

        cInt :: Integer
        cInt = case const of
          CIntConst (CInteger a _ _) _ -> a
          _ -> unexpected const "non integer initialization"

        cRat :: Rational
        cRat = case const of
          CIntConst (CInteger a _ _) _ -> fromIntegral a
          CFloatConst (CFloat a) _     -> fromIntegral (read a :: Integer)
          _ -> unexpected const "non numeric initialization"

      (Just (CInitExpr c _), Nothing) -> evalDecl' env (a, Nothing, Just c)

      (Nothing, Just e') -> do
        e <- evalExpr env e'
        v <- if isVolatile typInfo
          then return $ Volatile name typ $ posOf n
          else do
            m <- get
            put m { nextId = nextId m + 1 }
            return $ Local name typ (nextId m) (posOf e')
        assign True (posOf e') v e
        addVar env v
      _ -> notSupported i "variable declaration"
    _ -> notSupported d "arrays, pointers, or functional pointers (So what good is this tool anyway?)"
      

evalFunc :: Env -> CFunDef -> M Env
evalFunc env f =  do
  when (typ /= Void || not (null args)) $ notSupported f "non void f() function (How lame is this?)"
  return env { values = EnvFunction name (Function (length args) func) : values env }
  where
  (specs, (Ident name _ _), args, stat) = functionInfo f
  (_, typ) = typeInfo specs
  func _ = do --XXX
    evalStat env stat
    return false






-- | Assgins a value to a variable.
assign :: Bool -> Position -> V -> E -> M ()
assign decl n v a = case v of
  Volatile _ _ _ -> return ()
  _ -> do
    m <- get
    newAction $ Assign v $ if decl then a else Mux (enabled m) a (Var v) n  --XXX Is this correct for local and tmp variables?

-- | Latch a value at a point in time and return the expression of the latch variable.
latch :: Position -> E -> M E
latch _ (Var a)   = return (Var a)
latch _ (Const a) = return (Const a)
latch n a = do
  m <- get
  let v = Tmp (typeOf a) (nextId m) n
  put m { nextId = nextId m + 1 }
  assign True n v a
  return $ Var v

-- | Logical implication.  Assumes Bool inputs.
imply :: E -> E -> Position -> E
imply a b n = Or (Not a n) b n






-- | Extract relavent info from a function declaration.
functionInfo :: CFunDef -> ([CDeclSpec], Ident, [CDecl], CStat)
functionInfo (CFunDef specs (CDeclr (Just ident) [(CFunDeclr (Right (args, False)) _ _)] Nothing [] _) [] stmt _) = (specs, ident, args, stmt)
functionInfo f = notSupported f "function"




-- | Topologically sorts functions based on dependencies.
sortFunctions :: [CFunDef] -> [CFunDef]
sortFunctions fs = case topo $ map functionDeps fs of
  Left a -> notSupported none $ "recursive functions somewhere among: " ++ show a
  Right a -> [ f | a <- a, f <- fs, functionName f == a ]

topo :: Eq a => [(a, [a])] -> Either [a] [a]
topo a = topo' [] a
  where
  topo' a [] = Right a
  topo' done waiting = if null next then Left $ fst $ unzip waiting else topo' (done ++ next) stillWaiting
    where
    next = [ a | (a, deps) <- waiting, all (flip elem done) deps ]
    stillWaiting = [ (a, deps) | (a, deps) <- waiting, notElem a next ]


functionName :: CFunDef -> String
functionName f = name where (_, Ident name _ _, _, _) = functionInfo f

afvFunctionNames :: [String]
afvFunctionNames = ["assert", "assume"]

-- | Analyzes a function for dependencies.
functionDeps :: CFunDef -> (String, [String])
functionDeps f = (functionName f, filter (flip notElem afvFunctionNames) $ nub $ fb [] [CNestedFunDef f])
  where
  rewrite :: CFunDef -> CStat
  rewrite f = CCompound [] (map CBlockDecl args ++ [CBlockStmt stat]) none where (_, _, args, stat) = functionInfo f
  fs :: [String] -> CStat -> [String]
  fs env a = case a of
    CLabel _ a _ _ -> fs env a
    CCase a b _ -> fe env a ++ fs env b
    CCases a b c _ -> fe env a ++ fe env b ++ fs env c
    CDefault a _ -> fs env a
    CExpr (Just a ) _ -> fe env a
    CExpr Nothing _ -> []
    CCompound _ a _ -> fb env a
    CIf a b Nothing _ -> fe env a ++ fs env b
    CIf a b (Just c) _ -> fe env a ++ fs env b ++ fs env c
    CSwitch a b _ -> fe env a ++ fs env b
    CWhile a b _ _ -> fe env a ++ fs env b
    CFor _ _ _ _ _ -> notSupported a "for loops"
    CGoto _ _ -> []
    CGotoPtr a _ -> fe env a
    CCont _ -> []
    CBreak _ -> []
    CReturn Nothing _ -> []
    CReturn (Just a) _ -> fe env a
    CAsm _ _ -> []
  fe :: [String] -> CExpr -> [String]
  fe env a = case a of
    CComma a _ -> fe' a
    CAssign _ a b _ -> fe' [a, b]
    CCond a (Just b) c _ ->  fe' [a, b, c]
    CCond a Nothing b _ ->  fe' [a, b]
    CBinary _ a b _ -> fe' [a, b]
    CCast _ a _ -> fe env a --XXX does not check cast declaration.
    CUnary _ a _ -> fe env a
    --CSizeofExpr CExpr NodeInfo
    --CSizeofType CDecl NodeInfo
    --CAlignofExpr CExpr NodeInfo
    --CAlignofType CDecl NodeInfo
    --CComplexReal CExpr NodeInfo
    --CComplexImag CExpr NodeInfo
    CIndex a b _ -> fe' [a, b]
    CCall (CVar (Ident n _ _) _) args _ -> (if elem n env then [] else [n]) ++ fe' args
    CCall _ _ _ -> notSupported a "non-named function references"
    CMember a _ _ _ -> fe env a
    CVar _ _ -> []
    CConst _ -> []
    --CCompoundLit CDecl CInitList NodeInfo
    CStatExpr a _ -> fs env a
    CLabAddrExpr _ _ -> []
    --CBuiltinExpr CBuiltin
    _ -> notSupported a "expression"
    where
    fe' = concatMap $ fe env
  fb :: [String] -> [CBlockItem] -> [String]
  fb _ [] = []
  fb env (a:b) = case a of
    CBlockStmt a -> fs env a ++ fb env b
    CBlockDecl _ -> fb env b
    CNestedFunDef f -> fs env' (rewrite f) ++ fb env' b where env' = functionName f : env