{-# LANGUAGE TupleSections #-}

module Code where

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

data Datum
  = Atom Int -- unifies a constant
  | Struct Id -- unifies a structure with arity
  | VoidRef -- unifies with anything
  | LocalRef Int 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 -- something unifiable
  | NoGoal -- trivial goal (directly after head)
  | Invoke Builtin -- also directly after head
  | Goal -- a new goal (set head)
  | Call -- all seems okay, call the head's hoal
  | LastCall -- tail call the head's goal
  | Cut -- remove all alternative clauses of the current goal
  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, 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
    , heap :: Heap -- a snapshot of the heap (there's no trail; we rely on CoW copies in other choicepoints)
    , stk :: [(Code, Scope, [Cho])] -- remaining goals together with their vars and cuts
    , 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
    , strtable :: StrTable -- string table
    }
  deriving (Show)

type PrlgEnv a = StateT Interp (InputT IO) a

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

data Builtin =
  Builtin BuiltinFn

instance Show Builtin where
  show _ = "Builtin _"

codeStruct ::
     Monad m
  => (Datum -> m a)
  -> (Datum -> [a] -> m a)
  -> (Datum -> m (Either Int a))
  -> (Datum -> Int -> m a)
  -> m a
  -> Heap
  -> Code
  -> m (Code, a)
codeStruct atom struct local rec end heap = go
  where
    go [] = ([], ) <$> end
    go (U lr@(LocalRef _ _):cs) = do
      x <- local lr
      case x of
        Left ref -> (cs, ) <$> heapStruct atom struct rec heap ref
        Right a -> pure (cs, a)
    go (U s@(Struct (IR.Id _ arity)):cs) = eat arity cs >>= traverse (struct s)
    go (U x:cs) = (cs, ) <$> atom x
    go cs = (cs, ) <$> end
    eat n cs
      | n <= 0 = pure (cs, [])
      | otherwise = do
        (rest, a) <- go cs
        fmap (a :) <$> eat (n - 1) rest

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