-- | Symbolic STG Execution Support Architecture
module SSTG.Core.Execution.Support
    ( SymbolicT(..)
    , fmap
    , pure
    , (<*>)
    , return
    , (>>=)

    , State(..)
    , Symbol(..)
    , Status(..)
    , Stack
    , Frame(..)
    , MemAddr
    , Value(..)
    , Locals
    , Heap
    , HeapObj(..)
    , Globals
    , Code(..)
    , PathCons
    , Constraint(..)
    , SymLinks

    , nameOccStr
    , nameUnique
    , varName
    , nullAddr
    , addrInt

    , emptyStack
    , popStack
    , pushStack
    , stackToList

    , emptyLocals
    , lookupLocals
    , insertLocals
    , insertLocalsList
    , localsToList

    , emptyHeap
    , lookupHeap
    , allocHeap
    , allocHeapList
    , insertHeap
    , insertHeapList
    , heapToList

    , emptyGlobals
    , lookupGlobals
    , insertGlobals
    , insertGlobalsList
    , globalsToList

    , emptyPathCons
    , insertPathCons
    , insertPathConsList
    , pathconsToList

    , emptySymLinks
    , insertSymLinks
    , symlinksToList

    , lookupValue
    , vlookupHeap
    , memAddrType
    ) 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
             | AddrObj MemAddr
             | 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 are the conjunctive normal form of `Constraint`s.
newtype PathCons = PathCons [Constraint] deriving (Show, Eq, Read)

-- | Constraints denote logical paths taken in program execution thus far.
data Constraint = Constraint (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` imique `Int`.
nameUnique :: Name -> Int
nameUnique (Name _ _ _ unq) = unq

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

-- | Null `MemAddr`.
nullAddr :: MemAddr
nullAddr = MemAddr 0

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

-- | Empty `Stack.
emptyStack :: Stack
emptyStack = Stack []

-- | `Stack` pop.
popStack :: Stack -> Maybe (Frame, Stack)
popStack (Stack [])     = Nothing
popStack (Stack (f:fs)) = Just (f, Stack fs)

-- | `Stack` push.
pushStack :: Frame -> Stack -> Stack
pushStack frame (Stack frames) = Stack (frame : frames)

-- | `Stack` as list of `Frame`s.
stackToList :: Stack -> [Frame]
stackToList (Stack frames) = frames

-- | Empty `Locals`.
emptyLocals :: Locals
emptyLocals = Locals M.empty

-- | `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

-- | `Locals` to key value pairs.
localsToList :: Locals -> [(Name, Value)]
localsToList (Locals lmap) = M.toList lmap

-- | Empty `Heap`.
emptyHeap :: Heap
emptyHeap = Heap M.empty (MemAddr 0)

-- | `Heap` lookup.
lookupHeap :: MemAddr -> Heap -> Maybe HeapObj
lookupHeap addr (Heap hmap prev) = case M.lookup addr hmap of
    Just (AddrObj redir) -> lookupHeap redir (Heap hmap prev)
    mb_hobj              -> mb_hobj

-- | `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 ((addrInt 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

-- | `Heap` to key value pairs.
heapToList :: Heap -> [(MemAddr, HeapObj)]
heapToList (Heap hmap _) = M.toList hmap

-- | Empty `Globals`.
emptyGlobals :: Globals
emptyGlobals = Globals M.empty

-- | `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

-- | `Globals` to key value pairs.
globalsToList :: Globals -> [(Name, Value)]
globalsToList (Globals gmap) = M.toList gmap

-- | Empty `PathCons`.
emptyPathCons :: PathCons
emptyPathCons = PathCons []

-- | `PathCons` insertion.
insertPathCons :: Constraint -> PathCons -> PathCons
insertPathCons cons (PathCons pcs) = PathCons (cons : pcs)

-- | Insert a list of `Constraint`s into a `PathCons`.
insertPathConsList :: [Constraint] -> PathCons -> PathCons
insertPathConsList conss pathcons = foldr insertPathCons pathcons conss

-- | `PathCons` to list of `Constraint`s.
pathconsToList :: PathCons -> [Constraint]
pathconsToList (PathCons conss) = conss

-- | Empty `SymLinks`.
emptySymLinks :: SymLinks
emptySymLinks = SymLinks M.empty

-- | `SymLinks` insertion.
insertSymLinks :: Name -> Name -> SymLinks -> SymLinks
insertSymLinks old new (SymLinks smap) = SymLinks (M.insert old new smap)

-- | `SymLinks` to list of key value pairs.
symlinksToList :: SymLinks -> [(Name, Name)]
symlinksToList (SymLinks smap) = M.toList smap

--   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
    Just val -> Just 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
    case hobj of
        AddrObj redir       -> memAddrType redir heap
        LitObj lit          -> Just (litType lit)
        SymObj (Symbol s _) -> Just (varType s)
        ConObj dcon _       -> Just (dataConType dcon)
        FunObj ps e _       -> Just (foldr FunTy (exprType e) (map varType ps))
        Blackhole           -> Just Bottom