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

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

    , nameOccStr
    , nameUnique
    , varName
    , null_addr
    , addrInt

    , init_status
    , incStatusSteps
    , updateStatusId

    , empty_stack
    , popStack
    , pushStack
    , stackToList

    , empty_locals
    , lookupLocals
    , insertLocals
    , insertLocalsList
    , localsToList

    , empty_heap
    , lookupHeap
    , allocHeap
    , allocHeapList
    , insertHeap
    , insertHeapList
    , heapToList

    , empty_globals
    , lookupGlobals
    , insertGlobals
    , insertGlobalsList
    , globalsToList

    , empty_pathcons
    , insertPathCons
    , insertPathConsList
    , pathconsToList

    , lookupValue
    , vlookupHeap
    , memAddrType
    ) where

import SSTG.Core.Language.Syntax
import SSTG.Core.Language.Typing

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

-- 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`.
null_addr :: MemAddr
null_addr = MemAddr 0

-- | `MemAddr`'s `Int` value.
addrInt :: MemAddr -> Int
addrInt (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 new_id status = status { status_id = new_id
                                      , 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:frames)) = Just (frame, Stack frames)

-- | `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`.
empty_locals :: Locals
empty_locals = 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 (k, v) (Locals lmap) = Locals (M.insert (varName k) v lmap)

-- | List insertion into `Locals`.
insertLocalsList :: [(Var, Value)] -> Locals -> Locals
insertLocalsList kvs locals = foldr insertLocals locals kvs

-- | `Locals` to key value pairs.
localsToList :: Locals -> [(Name, Value)]
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 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 (k, v) (Heap hmap prev) = Heap (M.insert k v hmap) prev

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

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

-- | Empty `Globals`.
empty_globals :: Globals
empty_globals = 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 (k, v) (Globals gmap) = Globals (M.insert (varName k) v 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 kvs globals = foldr insertGlobals globals kvs

-- | `Globals` to key value pairs.
globalsToList :: Globals -> [(Name, Value)]
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`.
insertPathConsList :: [Constraint] -> PathCons -> PathCons
insertPathConsList 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.

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