{- VAM 2P, done the lazy way -}
module Interpreter where

import Data.Function
import qualified Data.Map as M

data StrTable =
  StrTable Int (M.Map String Int) (M.Map Int String)
  deriving (Show)

emptystrtable = StrTable 0 M.empty M.empty

strtablize t@(StrTable nxt fwd rev) str =
  case fwd M.!? str of
    Just i -> (t, i)
    _ -> (StrTable (nxt + 1) (M.insert str nxt fwd) (M.insert nxt str rev), nxt)

data Id =
  Id
    { str :: Int
    , arity :: Int
    }
  deriving (Show, Eq, Ord)

data Datum
  = Atom Int -- unifies a constant
  | Struct Id -- unifies a structure with arity
  | VoidVar -- in code this unifies with anything; everywhere else this is an unbound variable
  | LocalRef Int -- local variable idx
  | HeapRef Int -- heap structure idx
  deriving (Show, Eq, Ord)

data Instr
  = U Datum -- something unifiable
  | NoGoal -- trivial goal
  | Goal -- we start a new goal, set up backtracking etc
  | Call -- all seems okay, call the goal
  | LastCall -- tail call the 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

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 (TODO can we externalize?)
    , cur :: Cho -- the choice that is being evaluated right now
    , cho :: [Cho] -- remaining choice points
    }
  deriving (Show)

prove :: Code -> Defs -> (Interp, Either String Bool)
prove g ds =
  let i0 =
        Interp
          { defs = ds
          , cur =
              Cho
                { hed = g
                , hvar = emptyScope
                , gol = [LastCall]
                , gvar = emptyScope
                , heap = emptyHeap
                , stk = []
                , cut = []
                }
          , cho = []
          }
      run (Left x) = x
      run (Right x) = run $ proveStep Right (\i e -> Left (i, e)) x
   in run (Right i0)

data Dereferenced
  = FreeRef Int
  | BoundRef Int Datum
  | NoRef

{- this gonna need Either String Bool for errors later -}
proveStep :: (Interp -> a) -> (Interp -> Either String Bool -> a) -> Interp -> a
proveStep c f i = go i
  where
    ifail msg = f i $ Left msg
    tailcut [LastCall] chos _ = Just chos
    tailcut [LastCall, Cut] _ cut = Just cut
    tailcut _ _ _ = Nothing
    withDef fn
      | Just d <- defs i M.!? fn = ($ d)
      | otherwise = const $ ifail $ "no definition: " ++ show fn
    {- Backtracking -}
    backtrack i@Interp {cho = chos}
      {- if available, restore the easiest choicepoint -}
      | (cho:chos) <- chos = c i {cur = cho, cho = chos}
      {- if there's no other choice, answer no -}
      | otherwise = f i $ Right False
    {- Unification -}
    go i@Interp {cur = cur@Cho { hed = U h:hs
                               , gol = U g:gs
                               , heap = heap@(Heap _ hmap)
                               }} = unify h g
        {- termination tools -}
      where
        uok = c i {cur = cur {hed = hs, gol = gs}}
        setHeap r x =
          c i {cur = cur {hed = hs, gol = gs, heap = writeHeap r x heap}}
        {- heap tools -}
        deref x =
          case hmap M.!? x of
            Just (HeapRef x') ->
              if x == x'
                then FreeRef x'
                else deref x'
            Just x' -> BoundRef x x'
            _ -> NoRef
        writeHeap addr x (Heap nxt m) = Heap nxt (M.adjust (const x) addr m)
        newHeapVar (Heap nxt m) =
          (nxt, Heap (nxt + 1) (M.insert nxt (HeapRef nxt) m))
        allocLocal scope reg cont
          | Just addr <- scope M.!? reg = cont scope heap addr
          | (addr, heap') <- newHeapVar heap =
            cont (M.insert reg addr scope) heap' addr
        newHeapStruct addr s@(Struct Id {arity = arity}) cont = undefined
        {- simple cases first -}
        unify VoidVar VoidVar = uok
        unify (Atom a) (Atom b)
          | a == b = uok
        unify VoidVar (Atom _) = uok
        unify (Atom _) VoidVar = uok
        unify (Struct a) (Struct b)
          | a == b = uok
        {- unifying a struct with void must cause us to skip the void -}
        unify VoidVar (Struct Id {arity = a}) =
          c i {cur = cur {hed = replicate a (U VoidVar) ++ hs, gol = gs}}
        unify (Struct Id {arity = a}) VoidVar =
          c i {cur = cur {hed = hs, gol = replicate a (U VoidVar) ++ gs}}
        {- handle local refs; first ignore their combination with voids to save memory -}
        unify (LocalRef _) VoidVar = uok
        unify VoidVar (LocalRef _) = uok
        {- allocate heap for LocalRefs and retry with HeapRefs -}
        unify (LocalRef hv) (LocalRef gv) = undefined -- avoid allocating 2 things
        unify (LocalRef hv) _ = undefined
        unify _ (LocalRef gv) = undefined
        {- handle heap refs; first ignore their combination with voids again -}
        unify (HeapRef _) VoidVar = uok
        unify VoidVar (HeapRef _) = uok
        {- actual HeapRefs, these are dereferenced and then unified; decide between copying and linking -}
        unify (HeapRef hr') g =
          case deref hr' of
            FreeRef hr ->
              case g of
                atom@(Atom _) -> setHeap hr atom
                s@(Struct _) ->
                  newHeapStruct
                    hr
                    s
                    (\nhs nheap ->
                       c i {cur = cur {hed = nhs ++ hs, gol = gs, heap = nheap}})
                HeapRef gr' ->
                  case deref gr' of
                    FreeRef gr -> setHeap hr (HeapRef gr)
                    BoundRef addr _ -> setHeap hr (HeapRef addr)
                    _ -> ifail "dangling goal ref (from head ref)"
            BoundRef _ atom@(Atom a) -> unify atom g
            BoundRef addr struct@(Struct Id {arity = arity}) ->
              c
                i
                  { cur =
                      cur
                        { hed =
                            U struct :
                            [U (HeapRef $ addr + i) | i <- [1 .. arity]] ++ hs
                        , gol = U g : gs
                        }
                  }
            _ -> ifail "dangling head ref"
        unify h (HeapRef gr') =
          case deref gr' of
            FreeRef gr ->
              case h of
                atom@(Atom _) -> setHeap gr atom
                s@(Struct _) ->
                  newHeapStruct
                    gr
                    s
                    (\ngs nheap ->
                       c i {cur = cur {hed = hs, gol = ngs ++ gs, heap = nheap}})
            BoundRef _ atom@(Atom b) -> unify h atom
            BoundRef addr struct@(Struct Id {arity = arity}) ->
              c
                i
                  { cur =
                      cur
                        { hed = U h : hs
                        , gol =
                            U struct :
                            [U (HeapRef $ addr + i) | i <- [1 .. arity]] ++ gs
                        }
                  }
            _ -> ifail "dangling goal ref"
        unify _ _ = backtrack i
    {- Resolution -}
    go i@Interp { cur = cur@Cho { hed = hed
                                , hvar = hvar
                                , gol = gol
                                , gvar = gvar
                                , heap = heap
                                , stk = stk
                                , cut = cut
                                }
                , cho = chos
                }
      {- top-level success -}
      | [NoGoal] <- hed
      , Just nchos <- tailcut gol chos cut
      , [] <- stk =
        f i {cur = cur {hed = [], gol = []}, cho = nchos} $ Right True
      {- cut before the first goal (this solves all cuts in head) -}
      | Cut:hs <- hed = c i {cur = cur {hed = hs}, cho = cut}
      {- succeed and return to caller -}
      | [NoGoal] <- hed
      , Just nchos <- tailcut gol chos cut
      , (Goal:U (Struct fn):gs, ngvar, _):ss <- stk =
        withDef fn $ \(hs:ohs) ->
          c
            i
              { cur =
                  cur
                    { hed = hs
                    , hvar = emptyScope
                    , gol = gs
                    , gvar = ngvar
                    , stk = ss
                    }
              , cho =
                  [Cho oh emptyScope gs ngvar heap ss nchos | oh <- ohs] ++
                  nchos
              }
      {- succeed and return to caller, and the caller wants a cut -}
      | [NoGoal] <- hed
      , Just _ <- tailcut gol chos cut
      , (Cut:Goal:U (Struct fn):gs, ngvar, rchos):ss <- stk =
        withDef fn $ \(hs:ohs) ->
          c
            i
              { cur =
                  cur
                    { hed = hs
                    , hvar = emptyScope
                    , gol = gs
                    , gvar = ngvar
                    , stk = ss
                    }
              , cho =
                  [Cho oh emptyScope gs ngvar heap ss rchos | oh <- ohs] ++
                  rchos
              }
      {- start matching next goal -}
      | [NoGoal] <- hed
      , (Call:Goal:U (Struct fn):gs) <- gol =
        withDef fn $ \(hs:ohs) ->
          c
            i
              { cur = cur {hed = hs, hvar = emptyScope, gol = gs}
              , cho =
                  [Cho oh emptyScope gs gvar heap stk chos | oh <- ohs] ++ chos
              }
      {- start matching next goal after a cut -}
      | [NoGoal] <- hed
      , (Call:Cut:Goal:U (Struct fn):gs) <- gol =
        withDef fn $ \(hs:ohs) ->
          c
            i
              { cur = cur {hed = hs, hvar = emptyScope, gol = gs}
              , cho =
                  [Cho oh emptyScope gs gvar heap stk cut | oh <- ohs] ++ cut
              }
      {- goal head matching succeeded, make a normal call -}
      | (Goal:U (Struct fn):ngs) <- hed
      , (Call:gs) <- gol =
        withDef fn $ \(hs:ohs) ->
          let nstk = (gs, gvar, chos) : stk
           in c i
                  { cur =
                      cur
                        { hed = hs
                        , hvar = emptyScope
                        , gol = ngs
                        , gvar = hvar
                        , stk = nstk
                        }
                  , cho =
                      [Cho oh emptyScope ngs hvar heap nstk chos | oh <- ohs] ++
                      chos
                  }
      {- successful match continued by tail call -}
      | (Goal:U (Struct fn):ngs) <- hed
      , Just nchos <- tailcut gol chos cut =
        withDef fn $ \(hs:ohs) ->
          c
            i
              { cur = cur {hed = hs, hvar = emptyScope, gol = ngs, gvar = hvar}
              , cho =
                  [Cho oh emptyScope ngs hvar heap stk nchos | oh <- ohs] ++
                  nchos
              }
    {- The End -}
    go _ = ifail "code broken: impossible instruction combo"