module Language.Elaborator (elaborate, ElabError(..)) where

import Language.AST
import qualified IR.Boolean as B
import qualified Data.Map.Strict as Map

data ElabError
  = UnboundVariable String
  | UndefinedFunction String
  | ArityMismatch String Int Int
  deriving (Show)

type Env = Map.Map String ([String], Expr)

elaborate :: Module -> Either ElabError [B.Circuit]
elaborate (Module _ stmts) = do
  let env = buildEnv stmts
  mapM (elabStmt env) (filter isVerifiable stmts)

buildEnv :: [Stmt] -> Env
buildEnv = foldl addDef Map.empty
  where
    addDef env (SDef (Ident name) params body) =
      Map.insert name (map identName params, body) env
    addDef env _ = env

isVerifiable :: Stmt -> Bool
isVerifiable (SAssert _) = True
isVerifiable (SProve _ _) = True
isVerifiable _ = False

elabStmt :: Env -> Stmt -> Either ElabError B.Circuit
elabStmt env (SAssert expr) = do
  bexpr <- elabExpr env Map.empty expr
  Right (B.Circuit "assert" [] bexpr)
elabStmt env (SProve (Ident name) expr) = do
  bexpr <- elabExpr env Map.empty expr
  Right (B.Circuit name [] bexpr)
elabStmt _ _ = Right (B.Circuit "_skip" [] B.BTrue)

elabExpr :: Env -> Map.Map String B.BExpr -> Expr -> Either ElabError B.BExpr
elabExpr _ locals (EVar (Ident v)) =
  case Map.lookup v locals of
    Just bv -> Right bv
    Nothing -> Left (UnboundVariable v)
elabExpr _ _ (ELit True) = Right B.BTrue
elabExpr _ _ (ELit False) = Right B.BFalse
elabExpr env locals (ENand l r) = do
  bl <- elabExpr env locals l
  br <- elabExpr env locals r
  Right (B.BNand bl br)
elabExpr env locals (EApp (Ident fname) args) =
  case Map.lookup fname env of
    Nothing -> Left (UndefinedFunction fname)
    Just (params, body) -> do
      if length params /= length args
        then Left (ArityMismatch fname (length params) (length args))
        else do
          bargs <- mapM (elabExpr env locals) args
          let newLocals = Map.union (Map.fromList (zip params bargs)) locals
          elabExpr env newLocals body