-- | Symbolic STG Execution Models
module SSTG.Core.Execution.Models
    ( module SSTG.Core.Execution.Models
    ) where

import SSTG.Core.Syntax

import qualified Data.Map as M

-- | Symbolic Transformation represents transformations applied to some
-- State`(s). This is useful in allowing us to transfer from different actions
-- within the engine.
newtype SymbolicT s a = SymbolicT { run :: s -> (s, a) }

-- | Functor instance of Symbolic Transformation.
instance Functor (SymbolicT s) where
    fmap f st = SymbolicT (\s0 -> let (s1, a1) = (run st) s0 in (s1, f a1))

-- | Applicative instance of Symbolic Transformation.
instance Applicative (SymbolicT s) where
    pure a    = SymbolicT (\s -> (s, a))
    sf <*> st = SymbolicT (\s0 -> let (s1, a1) = (run st) s0
                                      (s2, f2) = (run sf) s1 in (s2, f2 a1))

-- | Monad instance of Symbolic Transformation.
instance Monad (SymbolicT s) where
    return a  = pure a
    st >>= fs = SymbolicT (\s0 -> let (s1, a1) = (run st) s0
                                      (s2, a2) = (run (fs a1)) s1 in (s2, a2))

-- | `State` contains the information necessary to perform symbolic execution.
-- Eval/Apply graph reduction semantics are used.
data State = State { state_status  :: Status
                   , state_stack   :: Stack
                   , state_heap    :: Heap
                   , state_globals :: Globals
                   , state_code    :: Code
                   , state_names   :: [Name]
                   , state_paths   :: PathCons
                   , state_links   :: SymLinks
                   } deriving (Show, Eq, Read)

-- | Symbolic variables. The @Maybe (Expr, Locals)@ can be used to trace the
-- source from which the symbolic variable was generated. For instance, this is
-- useful during symbolic function application.
data Symbol = Symbol Var (Maybe (Expr, Locals)) deriving (Show, Eq, Read)

-- | State status.
data Status = Status { steps :: Int
                     } deriving (Show, Eq, Read)

-- | Execution stack used in graph reduction semnatics.
newtype Stack = Stack [Frame] deriving (Show, Eq, Read)

-- | Frames of a stack.
data Frame = CaseFrame   Var [Alt] Locals
           | ApplyFrame  [Atom]    Locals
           | UpdateFrame MemAddr
           deriving (Show, Eq, Read)

-- | Memory address for things on the `Heap`.
newtype MemAddr = MemAddr Int deriving (Show, Eq, Read, Ord)

-- | A `Value` is something that we aim to reduce our current expression down
-- into. `MemAddr` is a pointer to an object on the heap, such as `FunObj` or
-- `ConObj`, which are "returned" from expression evaluation in this form.
data Value = LitVal Lit
           | MemVal MemAddr
           deriving (Show, Eq, Read)

-- | Locals binds a `Var`'s `Name` to its some `Value`.
newtype Locals = Locals (M.Map Name Value) deriving (Show, Eq, Read)

-- | Heaps map `MemAddr` to `HeapObj`, while keeping track of the last address
-- that was allocated. This allows us to consistently allocate fresh addresses
-- on the `Heap`.
data Heap = Heap (M.Map MemAddr HeapObj) MemAddr deriving (Show, Eq, Read)

-- | Heap objects.
data HeapObj = LitObj Lit
             | SymObj Symbol
             | ConObj DataCon [Value]
             | FunObj [Var] Expr Locals
             | Blackhole
             deriving (Show, Eq, Read)

-- | Globals are statically loaded at the time when a `State` is loaded.
-- However, because uninterpreted / out-of-scope variables are made symbolic
-- at runtime, it can be modified during execution.
newtype Globals = Globals (M.Map Name Value) deriving (Show, Eq, Read)

-- | Evaluation of the current expression. We are either evaluating, or ready
-- to return with some `Value`.
data Code = Evaluate Expr Locals
          | Return   Value
          deriving (Show, Eq, Read)

-- | Path constraints.
type PathCons = [PathCond]

-- | Path conditions denote logical paths taken in program execution thus far.
data PathCond = PathCond (AltCon, [Var]) Expr Locals Bool
              deriving (Show, Eq, Read)

-- | Symbolic link tables helps keep track of what names went to what, what?
newtype SymLinks = SymLinks (M.Map Name Name) deriving (Show, Eq, Read)

--   Simple functions that require only the immediate data structure.

-- | A `Name`'s occurrence string.
nameOccStr :: Name -> String
nameOccStr (Name occ _ _ _) = occ

-- | Name's unique `Int` key.
nameUnique :: Name -> Int
nameUnique (Name _ _ _ unq) = unq

-- | Variable name.
varName :: Var -> Name
varName (Var name _) = name

-- | `MemAddr`'s `Int` value.
memAddrInt :: MemAddr -> Int
memAddrInt (MemAddr int) = int

-- | `Locals` lookup.
lookupLocals :: Var -> Locals -> Maybe Value
lookupLocals var (Locals lmap) = M.lookup (varName var) lmap

-- | `Locals` insertion.
insertLocals :: Var -> Value -> Locals -> Locals
insertLocals var val (Locals lmap) = Locals lmap'
  where lmap' = M.insert (varName var) val lmap

-- | List insertion into `Locals`.
insertLocalsList :: [(Var, Value)] -> Locals -> Locals
insertLocalsList []               locals = locals
insertLocalsList ((var, val):vvs) locals = insertLocalsList vvs locals'
  where locals' = insertLocals var val locals

-- | `Heap` lookup.
lookupHeap :: MemAddr -> Heap -> Maybe HeapObj
lookupHeap addr (Heap hmap _) = M.lookup addr hmap

-- | `Heap` allocation. Updates the last `MemAddr` kept in the `Heap`.
allocHeap :: HeapObj -> Heap -> (Heap, MemAddr)
allocHeap hobj (Heap hmap prev) = (Heap hmap' addr, addr)
  where addr  = MemAddr ((memAddrInt prev) + 1)
        hmap' = M.insert addr hobj hmap

-- | Allocate a list of `HeapObj` in a `Heap`, returning in the same order the
-- `MemAddr` at which they have been allocated at.
allocHeapList :: [HeapObj] -> Heap -> (Heap, [MemAddr])
allocHeapList []           heap = (heap, [])
allocHeapList (hobj:hobjs) heap = (heapf, addr : as)
  where (heap', addr) = allocHeap hobj heap
        (heapf, as)   = allocHeapList hobjs heap'

-- | `Heap` direct insertion at a specific `MemAddr`.
insertHeap :: MemAddr -> HeapObj -> Heap -> Heap
insertHeap addr hobj (Heap hmap prev) = Heap hmap' prev
  where hmap' = M.insert addr hobj hmap

-- | Insert a list of `HeapObj` at specified `MemAddr` locations.
insertHeapList :: [(MemAddr, HeapObj)] -> Heap -> Heap
insertHeapList []                 heap = heap
insertHeapList ((addr, hobj):ahs) heap = insertHeapList ahs heap'
  where heap' = insertHeap addr hobj heap

-- | `Globals` lookup.
lookupGlobals :: Var -> Globals -> Maybe Value
lookupGlobals var (Globals gmap) = M.lookup (varName var) gmap

-- | `Globals` insertion.
insertGlobals :: Var -> Value -> Globals -> Globals
insertGlobals var val (Globals gmap) = Globals gmap'
  where gmap' = M.insert (varName var) val gmap

-- | Insert a list of `Var` and `Value` pairs into `Globals`. This would
-- typically occur for new symbolic variables created from uninterpreted /
-- out-of-scope variables during runtime.
insertGlobalsList :: [(Var, Value)] -> Globals -> Globals
insertGlobalsList []               globals = globals
insertGlobalsList ((var, val):vvs) globals = insertGlobalsList vvs globals'
  where globals' = insertGlobals var val globals

--   Complex functions that involve multiple data structures.

-- | `Value` lookup from the `Locals` first, then `Globals`.
lookupValue :: Var -> Locals -> Globals -> Maybe Value
lookupValue var locals globals = case lookupLocals var locals of
    Nothing -> lookupGlobals var globals
    mb_val  -> mb_val

-- | `Heap` lookup. Returns the corresponding `MemAddr` and `HeapObj` if found.
vlookupHeap :: Var -> Locals -> Globals -> Heap -> Maybe (MemAddr, HeapObj)
vlookupHeap var locals globals heap = do
    val <- lookupValue var locals globals
    case val of
        LitVal _    -> Nothing
        MemVal addr -> lookupHeap addr heap >>= \hobj -> Just (addr, hobj)

-- | Type of `HeapObj` held at `MemAddr`, if found.
memAddrType :: MemAddr -> Heap -> Maybe Type
memAddrType addr heap = do
    hobj <- lookupHeap addr heap
    Just $ case hobj of
        Blackhole           -> Bottom
        LitObj lit          -> litType lit
        SymObj (Symbol s _) -> varType s
        ConObj dcon _       -> dataConType dcon
        FunObj prms expr _  -> foldr FunTy (exprType expr) (map varType prms)