-- | Rules
module SSTG.Core.Execution.Rules
    ( Rule(..)
    , reduce
    , isStateValueForm
    ) where

import SSTG.Core.Syntax
import SSTG.Core.Execution.Models
import SSTG.Core.Execution.Namer

import qualified Data.Map   as M

-- | Rules
data Rule = RuleAtomLit | RuleAtomLitPtr | RuleAtomValPtr | RuleAtomUnInt
          | RulePrimApp
          | RuleConApp
          | RuleFunAppExact | RuleFunAppUnder | RuleFunAppSym | RuleFunAppUnInt
          | RuleLet
          | RuleCaseLit | RuleCaseConPtr | RuleCaseAnyLit | RuleCaseAnyConPtr
                        | RuleCaseSym

          | RuleUpdateCThunk
          | RuleUpdateDLit | RuleUpdateDValPtr

          | RuleAltCCaseNonVal
          | RuleAltDLit | RuleAltDValPtr

          | RuleApplyCFunThunk  | RuleApplyCFunAppOver
          | RuleApplyDReturnFun | RuleApplyDReturnSym

          | RuleIdentity
          deriving (Show, Eq, Read, Ord)

--  Stack Independent Rules

-- | Is Heap Normal Form?
--   Does not include LitObj. i.e. if something points to this, nothing to do.
isHeapValueForm :: HeapObj -> Bool
isHeapValueForm (SymObj _)         = True
isHeapValueForm (ConObj _ _)       = True
isHeapValueForm (FunObj (_:_) _ _) = True
isHeapValueForm _                  = False

-- | Is Value Form
--   Either a lit or points to a heap value (not LitObj!)
isExprValueForm :: Expr -> Locals -> Globals -> Heap -> Bool
isExprValueForm (Atom (LitAtom _))   _      _       _    = True
isExprValueForm (Atom (VarAtom var)) locals globals heap =
    case vlookupHeap var locals globals heap of
        Just (_, hobj) -> isHeapValueForm hobj
        Nothing        -> False
isExprValueForm _ _ _ _ = False

-- | Is State Value?
isStateValueForm :: State -> Bool
isStateValueForm State { state_stack = stack
                       , state_heap  = heap
                       , state_code  = code }
  | Stack []          <- stack
  , Return (LitVal _) <- code = True

  | Stack []             <- stack
  , Return (MemVal addr) <- code
  , Just hobj            <- lookupHeap addr heap
  , isHeapValueForm hobj = True

  | otherwise = False

-- | Uneven Zipping
unevenZip :: [a] -> [b] -> ([(a, b)], Either [a] [b])
unevenZip as     []     = ([], Left as)
unevenZip []     bs     = ([], Right bs)
unevenZip (a:as) (b:bs) = ((a, b) : acc, rem)
  where (acc, rem) = unevenZip as bs

-- | Inject Type Closure
injTyClosure :: Type -> [Atom] -> Type
injTyClosure ty args = TyClosure ty (map atomType args)

-- | Bind Rhs to Heap Object
rhsToObj :: BindRhs -> Locals -> Globals -> Maybe HeapObj
rhsToObj (FunForm prms expr) locals _       = Just (FunObj prms expr locals)
rhsToObj (ConForm dcon args) locals globals = do
    arg_vals <- mapM (\a -> alookupValue a locals globals) args
    return (ConObj dcon arg_vals)

-- | Lift Let Binding
liftBinding :: Binding -> Locals -> Globals -> Heap -> Maybe (Heap, Locals)
liftBinding (Binding NonRec bnd) locals globals heap = do
    hobjs <- mapM (\r -> rhsToObj r locals globals) (map snd bnd)
    let (heap', addrs) = allocHeapList hobjs heap
    return (heap', locals)
liftBinding (Binding Rec bnd) (Locals lmap) globals heap = do
    let names    = map (varName . fst) bnd
    let hfakes   = map (const Blackhole) bnd
    -- Allocate dummy BlackholeS
    let (heap', addrs) = allocHeapList hfakes heap
    let mem_vals = map (\a -> MemVal a) addrs
    -- Use the allocated BlackholeS to construct the locals closure.
    let lmap'    = M.fromList (zip names mem_vals)
    let locals'  = Locals (M.union lmap' lmap)
    hobjs <- mapM (\r -> rhsToObj r locals' globals) (map snd bnd)
    let injects  = zip addrs hobjs
    return (insertHeapList injects heap', locals')

-- | Default Alts
defaultAlts :: [Alt] -> [Alt]
defaultAlts alts = [a | a @ (Alt Default _ _) <- alts]

-- | AltCon Based Alts
altConAlts :: [Alt] -> [Alt]
altConAlts alts = [a | a @ (Alt acon _ _) <- alts, acon /= Default]

-- | Match Lit Alts
matchLitAlts :: Lit -> [Alt] -> [Alt]
matchLitAlts lit alts = [a | a @ (Alt (LitAlt alit) _ _) <- alts, lit == alit]

-- | Match Data Alts
matchDataAlts :: DataCon -> [Alt] -> [Alt]
matchDataAlts dc alts = [a | a @ (Alt (DataAlt adc) _ _) <- alts, dc == adc]

-- | Negate Path Cons
negatePathCons :: PathCons -> PathCons
negatePathCons pcs = map (\(PathCond a e l b) -> (PathCond a e l (not b))) pcs

-- | Lift Sym Alt
liftSymAlt :: Var -> MemAddr -> Var -> Locals -> Heap -> [Name] -> Alt ->
              (Expr, Locals, Heap, PathCons, [Name])
liftSymAlt mvar addr cvar locals heap confs (Alt ac params expr) =
    (expr, locals', heap', pcons, confs')
  where pre_names = freshSeededNameList (map varName params) confs
        sym_vars  = map (\(p, n) -> Var n (varType p)) (zip params pre_names)
        sym_objs  = map (\v -> SymObj (Symbol v)) sym_vars
        (heap', addrs) = allocHeapList sym_objs heap
        mem_vals  = map (\a -> MemVal a) addrs
        llist     = (cvar, MemVal addr) : zip params mem_vals
        locals'   = insertLocalsList llist locals
        mxpr      = Atom (VarAtom mvar)
        pcons     = [PathCond (Alt ac params expr) mxpr locals' True]
        confs'    = pre_names ++ confs

-- | Alt Closure to State
altClosureToState :: State -> (Expr, Locals, Heap, PathCons, [Name]) -> State
altClosureToState state (expr, locals, heap, pcons, confs) = state'
  where state' = state { state_heap  = heap
                       , state_code  = Evaluate expr locals
                       , state_names = confs ++ state_names state
                       , state_paths = pcons ++ state_paths state }

-- | Reduce
reduce :: State -> Maybe (Rule, [State])
reduce state @ State { state_stack   = stack
                     , state_heap    = heap
                     , state_globals = globals
                     , state_code    = code
                     , state_names   = confs
                     , state_paths   = paths }

  -- Stack Independent Rules

  -- Atom Lit
  | Evaluate (Atom (LitAtom lit)) _ <- code = do
    return ( RuleAtomLit
           , [state { state_code = Return (LitVal lit) }])

  -- Atom Lit Pointer
  | Evaluate (Atom (VarAtom var)) locals <- code
  , Just (_, hobj) <- vlookupHeap var locals globals heap
  , LitObj lit     <- hobj = do
    return ( RuleAtomLitPtr
           , [state { state_code = Evaluate (Atom (LitAtom lit)) locals }])

  -- Rule Atom Val Pointer
  | Evaluate (Atom (VarAtom var)) locals <- code
  , Just (addr, hobj) <- vlookupHeap var locals globals heap
  , isHeapValueForm hobj = do
    return ( RuleAtomValPtr
           , [state { state_code = Return (MemVal addr) }])

  -- Rule Atom Uninterpreted
  | Evaluate (Atom (VarAtom uvar)) locals <- code
  , Nothing <- vlookupHeap uvar locals globals heap = do
    let sname    = freshSeededName (varName uvar) confs
    let svar     = Var sname (varType uvar)
    let sym      = Symbol svar
    let (heap', addr) = allocHeap (SymObj sym) heap
    let globals' = insertGlobals uvar (MemVal addr) globals
    return ( RuleAtomUnInt
           , [state { state_heap    = heap'
                    , state_globals = globals'
                    , state_code    = Evaluate (Atom (VarAtom uvar)) locals
                    , state_names   = sname : confs }])

  -- Prim Function App
  | Evaluate (PrimApp pfun args) locals <- code = do
    let eval = SymLitEval pfun args
    return ( RulePrimApp
           , [state { state_code = Evaluate (Atom (LitAtom eval)) locals }])

  -- | Rule Con App
  | Evaluate (ConApp dcon args) locals <- code = do
    arg_vals <- mapM (\a -> alookupValue a locals globals) args
    let (heap', addr) = allocHeap (ConObj dcon arg_vals) heap
    return ( RuleConApp
           , [state { state_heap = heap'
                    , state_code = Return (MemVal addr) }])

  -- | Rule Fun App Exact
  | Evaluate (FunApp fun args) locals <- code
  , Just (_, hobj)              <- vlookupHeap fun locals globals heap
  , FunObj params expr fun_locs <- hobj
  , length params == length args = do
    arg_vals <- mapM (\a -> alookupValue a locals globals) args
    let fun_locs' = insertLocalsList (zip params arg_vals) fun_locs
    return ( RuleFunAppExact
           , [state { state_code = Evaluate expr fun_locs' }])

  -- Rule Fun App Under
  | Evaluate (FunApp fun args) locals <- code
  , Just (_, hobj)              <- vlookupHeap fun locals globals heap
  , FunObj params expr fun_locs <- hobj
  , (_, Left ex_prms)           <- unevenZip params args = do
    -- Set up existing closure first.
    arg_vals <- mapM (\a -> alookupValue a locals globals) args
    let fun_locs' = insertLocalsList (zip params arg_vals) fun_locs
    -- New Fun Object.
    let pfobj     = FunObj ex_prms expr fun_locs'
    let (heap', pfaddr) = allocHeap pfobj heap
    return ( RuleFunAppUnder
           , [state { state_heap = heap'
                    , state_code = Return (MemVal pfaddr) }])

  -- Rule Fun App Symbolic
  | Evaluate (FunApp sfun args) locals <- code
  , Just (_, hobj)       <- vlookupHeap sfun locals globals heap
  , SymObj (Symbol svar) <- hobj = do
    let sname = freshSeededName (varName svar) confs
    let svar' = Var sname (injTyClosure (varType svar) args)
    let sym   = Symbol svar'
    let (heap', addr) = allocHeap (SymObj sym) heap
    return ( RuleFunAppSym
           , [state { state_heap  = heap'
                    , state_code  = Return (MemVal addr)
                    , state_names = sname : confs }])

  -- Rule Fun App Uninterpreted
  | Evaluate (FunApp ufun args) locals <- code
  , Nothing  <- vlookupHeap ufun locals globals heap = do
    let sname    = freshSeededName (varName ufun) confs
    let svar     = Var sname (varType ufun)
    let sym      = Symbol svar
    let (heap', addr) = allocHeap (SymObj sym) heap
    let globals' = insertGlobals ufun (MemVal addr) globals
    return ( RuleFunAppUnInt
           , [state { state_heap    = heap'
                    , state_globals = globals'
                    , state_code    = Evaluate (FunApp ufun args) locals
                    , state_names   = sname : confs }])

  -- Rule Let
  | Evaluate (Let bnd expr) locals <- code = do
    (heap', locals') <- liftBinding bnd locals globals heap
    return ( RuleLet
           , [state { state_heap = heap'
                    , state_code = Evaluate expr locals' }])

  -- Rule Case Lit
  | Evaluate (Case (Atom (LitAtom lit)) cvar alts) locals <- code
  , (Alt _ _ expr):_ <- matchLitAlts lit alts = do
    -- Account for cvar.
    let locals' = insertLocals cvar (LitVal lit) locals
    return ( RuleCaseLit
           , [state { state_code = Evaluate expr locals' }])

  -- Rule Case Con Pointer
  | Evaluate (Case (Atom (VarAtom mvar)) cvar alts) locals <- code
  , Just (addr, hobj)     <- vlookupHeap mvar locals globals heap
  , ConObj dcon vals      <- hobj
  , (Alt _ params expr):_ <- matchDataAlts dcon alts
  , length params == length vals = do
    -- Account for cvar.
    let llist   = (cvar, MemVal addr) : zip params vals
    let locals' = insertLocalsList llist locals
    return ( RuleCaseConPtr
           , [state { state_code = Evaluate expr locals' }])

  -- Rule Case Any Lit
  | Evaluate (Case (Atom (LitAtom lit)) cvar alts) locals <- code
  , []               <- matchLitAlts lit alts
  , (Alt _ _ expr):_ <- defaultAlts alts = do
    -- Account for cvar.
    let locals' = insertLocals cvar (LitVal lit) locals
    return ( RuleCaseAnyLit
           , [state { state_code = Evaluate expr locals' }])

  -- Rule Case Any Con Pointer
  | Evaluate (Case (Atom (VarAtom mvar)) cvar alts) locals <- code
  , Just (addr, hobj) <- vlookupHeap mvar locals globals heap
  , ConObj dcon _     <- hobj
  , []                <- matchDataAlts dcon alts
  , (Alt _ _ expr):_  <- defaultAlts alts = do
    -- Account for cvar.
    let llist   = (cvar, (MemVal addr)) : []
    let locals' = insertLocalsList llist locals
    return ( RuleCaseAnyConPtr
           , [state { state_code = Evaluate expr locals' }])

  -- Rule Case Sym
  | Evaluate (Case (Atom (VarAtom mvar)) cvar alts) locals <- code
  , Just (addr, hobj)     <- vlookupHeap mvar locals globals heap
  , SymObj _              <- hobj
  , (acon_alts, def_alts) <- (altConAlts alts, defaultAlts alts)
  , length (acon_alts ++ def_alts) > 0 = do
    -- Remember to account for cvar.
    let acon_clss = map (liftSymAlt mvar addr cvar locals heap confs) acon_alts
    let def_clss  = map (liftSymAlt mvar addr cvar locals heap confs) def_alts
    -- Make AltCon states first.
    let acon_sts  = map (altClosureToState state) acon_clss
    -- Make Default states next.
    let all_pcons = concatMap (\(_, _, _, pc, _) -> pc) acon_clss
    let neg_pcons = negatePathCons all_pcons
    let def_clss' = map (\(e, l, h, p, c) -> (e, l, h, neg_pcons, c)) def_clss
    let def_sts   = map (altClosureToState state) def_clss'
    return (RuleCaseSym, acon_sts ++ def_sts)

  -- Stack Dependent Rules

  -- Rule Update Frame Create Thunk
  | Stack frames                         <- stack
  , Evaluate (Atom (VarAtom var)) locals <- code
  , Just (addr, hobj)       <- vlookupHeap var locals globals heap
  , FunObj [] expr fun_locs <- hobj = do  -- Thunk form.
    return ( RuleUpdateCThunk
           , [state { state_stack = Stack (UpdateFrame addr : frames)
                    , state_heap  = insertHeap addr Blackhole heap
                    , state_code  = Evaluate expr fun_locs }])

  -- Rule Update Frame Delete Lit
  | Stack (UpdateFrame frm_addr : rest) <- stack
  , Return (LitVal lit)                 <- code = do
    return ( RuleUpdateDLit
           , [state { state_stack = Stack rest
                    , state_heap  = insertHeap frm_addr (LitObj lit) heap
                    , state_code  = Return (LitVal lit) }])

  -- Rule Update Frame Delete Val Pointer
  | Stack (UpdateFrame frm_addr : rest) <- stack
  , Return (MemVal addr)                <- code
  , Just hobj <- lookupHeap addr heap
  , isHeapValueForm hobj = do
    return ( RuleUpdateDValPtr
           , [state { state_stack = Stack rest
                    , state_heap  = insertHeap frm_addr hobj heap
                    , state_code  = Return (MemVal addr) }])

  -- Rule Alt Frame Create Case Non LitVal or MemVal
  | Stack frames                          <- stack
  , Evaluate (Case mxpr cvar alts) locals <- code
  , not (isExprValueForm mxpr locals globals heap) = do
    return ( RuleAltCCaseNonVal
           , [state { state_stack = Stack (AltFrame cvar alts locals : frames)
                    , state_code  = Evaluate mxpr locals }])

    -- Rule Alt Frame Delete Lit
  | Stack (AltFrame cvar alts frm_locs : rest) <- stack
  , Return (LitVal lit)                        <- code = do
    let mxpr = Atom (LitAtom lit)
    return ( RuleAltDLit
           , [state { state_stack = Stack rest
                    , state_code  = Evaluate (Case mxpr cvar alts) frm_locs }])

  -- Rule Alt Frame Delete Heap Value
  | Stack (AltFrame cvar alts frm_locs : rest) <- stack
  , Return (MemVal addr)                       <- code
  , Just hobj <- lookupHeap addr heap
  , isHeapValueForm hobj = do
    let vname     = freshSeededName (varName cvar) confs
    let vvar      = Var vname (varType cvar)
    let mxpr      = Atom (VarAtom vvar)
    let frm_locs' = insertLocals vvar (MemVal addr) frm_locs
    return ( RuleAltDValPtr
           , [state { state_stack = Stack rest
                    , state_code  = Evaluate (Case mxpr cvar alts) frm_locs'
                    , state_names = vname : confs }])

  -- Rule Apply Frame Create Function Thunk
  | Stack frames                      <- stack
  , Evaluate (FunApp fun args) locals <- code
  , Just (_, hobj)          <- vlookupHeap fun locals globals heap
  , FunObj [] expr fun_locs <- hobj = do
    return ( RuleApplyCFunThunk
           , [state { state_stack = Stack (ApplyFrame args locals : frames)
                    , state_code  = Evaluate expr fun_locs }])

  -- Rule Apply Frame Create Function Over Application
  | Stack frames                      <- stack
  , Evaluate (FunApp fun args) locals <- code
  , Just (_, hobj)              <- vlookupHeap fun locals globals heap
  , FunObj params expr fun_locs <- hobj
  , (_, Right ex_args)          <- unevenZip params args = do
    arg_vals <- mapM (\a -> alookupValue a locals globals) args
    let fun_locs' = insertLocalsList (zip params arg_vals) fun_locs
    return ( RuleApplyCFunAppOver
           , [state { state_stack = Stack (ApplyFrame ex_args locals : frames)
                    , state_code  = Evaluate expr fun_locs' }])

  -- Rule Apply Frame Delete ReturnPtr Function
  | Stack (ApplyFrame args frm_locs : rest) <- stack
  , Return (MemVal addr)                    <- code
  , Just hobj    <- lookupHeap addr heap
  , FunObj _ _ _ <- hobj = do
    ftype <- memAddrType addr heap
    let fname     = freshName confs
    let fvar      = Var fname ftype
    let frm_locs' = insertLocals fvar (MemVal addr) frm_locs
    return ( RuleApplyDReturnFun
           , [state { state_stack = Stack rest
                    , state_code  = Evaluate (FunApp fvar args) frm_locs'
                    , state_names = fname : confs }])

  -- Rule Apply Frame Delete ReturnPtr Sym
  | Stack (ApplyFrame args frm_locs : rest) <- stack
  , Return (MemVal addr)                    <- code
  , Just hobj           <- lookupHeap addr heap
  , SymObj (Symbol sym) <- hobj = do
    let sname     = freshSeededName (varName sym) confs
    let svar      = Var sname (varType sym)
    let frm_locs' = insertLocals svar (MemVal addr) frm_locs
    return ( RuleApplyDReturnSym
           , [state { state_stack = Stack rest
                    , state_code  = Evaluate (FunApp svar args) frm_locs'
                    , state_names = sname : confs }])

  -- State is Value Form
  | isStateValueForm state = return (RuleIdentity, [state])

  -- Everything Broke!!!
  | otherwise = Nothing