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

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

    , null_addr
    , memAddrInt

    , init_status
    , incStatusSteps
    , updateStatusId

    , empty_stack
    , popStack
    , pushStack
    , stackToList

    , empty_locals
    , lookupLocals
    , insertLocalsVal
    , insertLocalsVals
    , localsToList

    , empty_heap
    , lookupHeap
    , insertHeapObj
    , insertHeapObjs
    , insertHeapRedir
    , allocHeapObj
    , allocHeapObjs
    , heapToList

    , empty_globals
    , lookupGlobals
    , insertGlobalsVal
    , insertGlobalsVals
    , globalsToList

    , empty_pathcons
    , insertPathCons
    , insertPathConss
    , pathConsToList

    , lookupVal
    , vlookupHeap
    , memAddrType
    ) where

import SSTG.Core.Language

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_path :: !PathCons
                   } 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 { status_id :: !Int
                     , status_parent :: !Int
                     , 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 `Val` 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 Val = LitVal Lit
         | MemVal MemAddr
         deriving (Show, Eq, Read)

-- | Locals binds a `Var`'s `Name` to its some `Val`.
newtype Locals = Locals (M.Map Name Val) 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 HeapRhs) MemAddr
          deriving (Show, Eq, Read)

-- | When we look up a `Heap`, we get an intermediary `MemAddr` for lookup
-- redirection, or just the `HeapObj` itself.
data HeapRhs = HeapRedir MemAddr
             | HeapObj HeapObj
             deriving (Show, Eq, Read)

-- | Heap objects.
data HeapObj = LitObj Lit
             | SymObj Symbol
             | ConObj DataCon [Val]
             | 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 Val) deriving (Show, Eq, Read)

-- | Evaluation of the current expression. We are either evaluating, or ready
-- to return with some `Val`.
data Code = Evaluate Expr Locals
          | Return Val
          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 Expr Locals Bool deriving (Show, Eq, Read)

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

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

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

-- | Initial `Status`.
init_status :: Status
init_status = Status { status_id = 1
                     , status_parent = 0
                     , status_steps = 0 }

-- | Increment `Status` steps.
incStatusSteps :: Status -> Status
incStatusSteps status = status { status_steps = (status_steps status) + 1 }

-- | Update the `Status` id.
updateStatusId :: Int -> Status -> Status
updateStatusId int status = status { status_id = int
                                   , status_parent = status_id status }

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

-- | `Stack` pop.
popStack :: Stack -> Maybe (Frame, Stack)
popStack (Stack []) = Nothing
popStack (Stack (frame:fs)) = Just (frame, 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

-- | `foldr` helper function that takes (A, B) into A -> B type inputs.
foldrPair :: (a -> b -> c -> c) -> (a, b) -> c -> c
foldrPair f (a, b) c = f a b c

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

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

-- | `Locals` insertion.
insertLocalsVal :: Var -> Val -> Locals -> Locals
insertLocalsVal k v (Locals lmap) = Locals (M.insert (varName k) v lmap)

-- | List insertion into `Locals`.
insertLocalsVals :: [(Var, Val)] -> Locals -> Locals
insertLocalsVals kvs locals = foldr (foldrPair insertLocalsVal) locals kvs

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

-- | Empty `Heap`.
empty_heap :: Heap
empty_heap = Heap M.empty null_addr

-- | `Heap` lookup.
lookupHeap :: MemAddr -> Heap -> Maybe HeapObj
lookupHeap addr (Heap hmap p) = case M.lookup addr hmap of
    Just (HeapRedir redir) -> lookupHeap redir (Heap hmap p)
    Just (HeapObj hobj) -> Just hobj
    Nothing -> Nothing

-- | `Heap` direct insertion at a specific `MemAddr`.
insertHeapObj :: MemAddr -> HeapObj -> Heap -> Heap
insertHeapObj k v (Heap hmap p) = Heap (M.insert k (HeapObj v) hmap) p

-- | Insert a list of `HeapObj` at specified `MemAddr` locations.
insertHeapObjs :: [(MemAddr, HeapObj)] -> Heap -> Heap
insertHeapObjs kvs heap = foldr (foldrPair insertHeapObj) heap kvs

-- | Insert a redirection `MemAddr` into the `Heap`.
insertHeapRedir :: MemAddr -> MemAddr -> Heap -> Heap
insertHeapRedir a r (Heap hmap p) = Heap (M.insert a (HeapRedir r) hmap) p

-- | `Heap` allocation. Updates the last `MemAddr` kept in the `Heap`.
allocHeapObj :: HeapObj -> Heap -> (Heap, MemAddr)
allocHeapObj hobj (Heap hmap p) = (heap', addr)
  where
    addr = MemAddr ((memAddrInt p) + 1)
    heap' = insertHeapObj addr hobj (Heap hmap addr)

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

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

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

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

-- | `Globals` insertion.
insertGlobalsVal :: Var -> Val -> Globals -> Globals
insertGlobalsVal k v (Globals gmap) = Globals (M.insert (varName k) v gmap)

-- | Insert a list of `Var` and `Val` pairs into `Globals`. This would
-- typically occur for new symbolic variables created from uninterpreted /
-- out-of-scope variables during runtime.
insertGlobalsVals :: [(Var, Val)] -> Globals -> Globals
insertGlobalsVals kvs globals = foldr (foldrPair insertGlobalsVal) globals kvs

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

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

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

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

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

-- Complex functions that involve multiple data structures.

-- | `Val` lookup from the `Locals` first, then `Globals`.
lookupVal :: Var -> Locals -> Globals -> Maybe Val
lookupVal var locals globals = case lookupLocals var locals of
    Just val -> Just val
    Nothing -> lookupGlobals var globals

-- | `Heap` lookup. Returns the appropriate. `MemAddr` and `HeapObj` if found.
vlookupHeap :: Var -> Locals -> Globals -> Heap -> Maybe (MemAddr, HeapObj)
vlookupHeap var locals globals heap = do
    val <- lookupVal 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
        LitObj lit -> Just (typeOf lit)
        SymObj (Symbol svar _) -> Just (typeOf svar)
        ConObj dcon _ -> Just (typeOf dcon)
        FunObj ps expr _ -> Just (foldr FunTy (typeOf expr) (map typeOf ps))
        Blackhole -> Just Bottom