{-# LANGUAGE TupleSections #-}

module Code where

import Control.Monad.Trans.State.Lazy (StateT)
import qualified Data.Map as M
import IR (Id(..), StrTable)
import Operators (Ops)
import Parser (PAST)
import System.Console.Haskeline (InputT)

data Datum
  = Atom Int -- unifies a symbolic constant
  | Number Int -- unifies a numeric constant
  | Struct Id -- unifies a structure with arity
  | VoidRef -- unifies with anything
  | LocalRef Int -- code-local variable idx (should never occur on heap)
  | HeapRef Int -- something further on the heap
  deriving (Show, Eq, Ord)

data Instr
  = U Datum -- unify/resolve something
  | Invoke Builtin -- give control to a builtin (invoked from head)
  | Done -- all done, can return
  | Cut -- remove choicepoints of the current goal
  | Choices [[Code]] -- split off several possibilities (push choicepoints)
  deriving (Show)

type Code = [Instr]

type Defs = M.Map Id [Code]

data Heap =
  Heap Int (M.Map Int Datum)
  deriving (Show)

emptyHeap = Heap 1 M.empty

type Scope = M.Map Int Int

emptyScope :: Scope
emptyScope = M.empty

data Cho =
  Cho
    { _hed :: Code -- head pointer
    , _hvar :: Scope -- variables unified in head (so far)
    , _gol :: Code -- goal pointer
    , _gvar :: Scope -- variables unified in the goal
    , _unis :: Int -- items left to unify
    , _retcut :: Bool -- cut after this goal succeeds
    , _heap :: Heap -- a snapshot of the heap (there's no trail; we rely on CoW copies in other choicepoints)
    , _stk :: [(Code, Scope, [Cho], Bool)] -- remaining goals together with their vars, cuts, and ret-cut flag
    , _cut :: [Cho] -- snapshot of choicepoints before entering
    }
  deriving (Show)

data Interp =
  Interp
    { _defs :: Defs -- global definitions for lookup
    , _cur :: Cho -- the choice that is being evaluated right now
    , _cho :: [Cho] -- remaining choice points
    , _ops :: Ops -- currently defined operators
    , _opstash :: [Ops] -- saved operators
    , _macrostash :: [(Maybe [Code], Maybe [Code])] -- saved expansion defs (load, query)
    , _strtable :: StrTable -- string table
    , _cmdq :: [(Bool, PAST)] -- isQuery, lexemes
    }
  deriving (Show)

type PrlgEnv = StateT Interp (InputT IO)

type InterpFn = PrlgEnv (Maybe (Either String Bool))

data Builtin =
  Builtin InterpFn

instance Show Builtin where
  show _ = "Builtin _"

data Dereferenced
  = FreeRef Int
  | BoundRef Int Datum
  | NoRef
  deriving (Show, Eq)

-- TRICKY: HeapRefs alone must not form a cycle other than the FreeRef case.
derefHeap :: Heap -> Int -> Dereferenced
derefHeap h@(Heap _ hmap) x =
  case hmap M.!? x of
    Just (HeapRef x') ->
      if x == x'
        then FreeRef x'
        else derefHeap h x'
    Just x' -> BoundRef x x'
    _ -> NoRef

writeHeap :: Int -> Datum -> Heap -> Heap
writeHeap addr x (Heap nxt m) = Heap nxt (M.adjust (const x) addr m)

newHeapVar :: Heap -> (Heap, Int)
newHeapVar heap = head <$> newHeapVars 1 heap

newHeapVars :: Int -> Heap -> (Heap, [Int])
newHeapVars n (Heap nxt m) =
  let addrs = [nxt + i - 1 | i <- [1 .. n]]
   in ( Heap (nxt + n) $
        foldr (uncurry M.insert) m [(a, HeapRef a) | a <- addrs]
      , addrs)

withNewHeapStruct :: Int -> Datum -> Heap -> ([Datum] -> Heap -> a) -> a
withNewHeapStruct addr s@(Struct Id {arity = arity}) heap cont =
  let (Heap nxt' m', addrs) = newHeapVars (arity + 1) heap
      m'' = M.insert addr (HeapRef $ head addrs) . M.insert (head addrs) s $ m'
   in cont (map HeapRef $ tail addrs) (Heap nxt' m'')

heapStruct ::
     Monad m
  => (Datum -> m a)
  -> (Datum -> [a] -> m a)
  -> (Datum -> Int -> m a)
  -> Heap
  -> Int
  -> m a
heapStruct atom struct rec (Heap _ heap) hr = go [hr] hr
  where
    go visited ref
      | rr@(HeapRef r) <- heap M.! ref =
        if r == ref || r `elem` visited
          then rec rr ref
          else go (r : visited) r
      | s@(Struct (IR.Id _ arity)) <- heap M.! ref =
        sequence [go (ref + i : visited) (ref + i) | i <- [1 .. arity]] >>=
        struct s
      | x <- heap M.! ref = atom x