{- pražský přehledný stroj -}
module Interpreter where

import Code
  ( Builtin(..)
  , Cho(..)
  , Code
  , Datum(..)
  , Dereferenced(..)
  , Instr(..)
  , InterpFn
  , derefHeap
  , emptyHeap
  , emptyScope
  , newHeapVar
  , withNewHeapStruct
  , writeHeap
  )
import CodeLens
import qualified Data.Map as M
import Env (PrlgEnv)
import IR (Id(..), StrTable(..))
import Lens.Family2.State

prove :: Code -> PrlgEnv (Either String Bool)
prove g = do
  cur .=
    Cho
      { _hed = g
      , _hvar = emptyScope
      , _gol = [Done]
      , _gvar = emptyScope
      , _unis = 0
      , _retcut = True
      , _heap = emptyHeap
      , _stk = []
      , _cut = []
      }
  cho .= []
  loop
  where
    loop = do
      x <- proveStep
      case x of
        Nothing -> loop -- not finished yet
        Just x -> return x

{- toplevel decision -}
proveStep :: InterpFn
proveStep = do
  u <- use (cur . unis)
  h <- use (cur . hed)
  case (u, h) of
    (0, []) -> goalStep
    (0, _) -> headStep h
    (_, _)
      | u > 0 -> unifyStep h
    _ -> err "invalid state"

err :: String -> InterpFn
err = return . Just . Left

{- toplevel choices -}
goalStep :: InterpFn
goalStep = do
  g <- use (cur . gol)
  case g of
    U (Struct _):gs -> undefined -- TODO these things NEED lens-family.
    [Done] -> undefined
    [Cut, Done] -> undefined
    Cut:gs -> undefined
    [Choices cs, Done] -> undefined
    [Choices cs, Cut, Done] -> undefined
    Choices cs:gs -> undefined
    _ -> err "invalid goal code"

headStep :: [Instr] -> InterpFn
headStep h = do
  g <- use (cur . gol)
  case (h, g) of
    ([Done], _) -> undefined
    ([Cut, Done], _) -> undefined
    (_, [Done]) -> undefined
    (_, [Cut, Done]) -> undefined
    (_, _) -> undefined

unifyStep h = do
  g <- use (cur . gol)
  case (h, g) of
    (U hd:_, U gd:_) -> undefined hd gd
    (_, _) -> err "invalid unification code"

{- helpers -}
backtrack :: InterpFn
backtrack = do
  chos <- use cho
  case chos of
    (c:cs)
      {- if available, restore the easiest choicepoint -}
     -> do
      cur .= c
      cho .= cs
      pure Nothing
    {- if there's no other choice, answer no -}
    _ -> pure . Just $ Right False

retCut :: InterpFn
retCut = undefined

cutHead :: InterpFn
cutHead = undefined

cutGoal :: InterpFn
cutGoal = undefined
{- original, TODO remove -}
{-proveStep :: InterpFn
proveStep = St.get >>= go
  where
    finish = pure . Just
    c i = St.put i >> pure Nothing
    ifail msg = finish $ Left msg
    tailcut [LastCall] chos _ = Just chos
    tailcut [LastCall, Cut] _ cut = Just cut
    tailcut _ _ _ = Nothing
    withDef fn cont = do
      d <- St.gets defs
      case d M.!? fn of
        Just d -> cont d
        _ -> do
          StrTable _ _ itos <- St.gets strtable
          ifail $
            "no definition: '" ++ (itos M.! str fn) ++ "'/" ++ show (arity fn)
    {- Unification -}
    go i@Interp {cur = cur@Cho {hed = U h:hs, gol = U g:gs, heap = heap}} =
      unify h g
      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 = derefHeap heap
        withNewLocal (LocalRef reg) scope cont
          | Just addr <- scope M.!? reg = cont scope heap addr
          | (heap', addr) <- newHeapVar heap =
            cont (M.insert reg addr scope) heap' addr
        {- simple cases first -}
        unify VoidRef VoidRef = uok
        unify (Atom a) (Atom b)
          | a == b = uok
        unify VoidRef (Atom _) = uok
        unify (Atom _) VoidRef = uok
        unify (Number a) (Number b)
          | a == b = uok
        unify VoidRef (Number _) = uok
        unify (Number _) VoidRef = uok
        unify (Struct a) (Struct b)
          | a == b = uok
        {- unifying a struct with void must cause us to skip the void -}
        unify VoidRef (Struct Id {arity = a}) =
          c i {cur = cur {hed = replicate a (U VoidRef) ++ hs, gol = gs}}
        unify (Struct Id {arity = a}) VoidRef =
          c i {cur = cur {hed = hs, gol = replicate a (U VoidRef) ++ gs}}
        {- handle local refs; first ignore their combination with voids to save memory -}
        unify (LocalRef _) VoidRef = uok -- TRICKY: builtins need to check if locals actually exist because of this
        unify VoidRef (LocalRef _) = uok
        {- allocate heap for LocalRefs and retry with HeapRefs -}
        unify lr@(LocalRef _) _ =
          withNewLocal lr (hvar cur) $ \hvar' heap' addr ->
            c
              i
                { cur =
                    cur
                      {hed = U (HeapRef addr) : hs, hvar = hvar', heap = heap'}
                }
        unify _ lr@(LocalRef _) =
          withNewLocal lr (gvar cur) $ \gvar' heap' addr ->
            c
              i
                { cur =
                    cur
                      {gol = U (HeapRef addr) : gs, gvar = gvar', heap = heap'}
                }
        {- handle heap refs; first ignore their combination with voids again -}
        unify (HeapRef _) VoidRef = uok
        unify VoidRef (HeapRef _) = uok
        {- actual HeapRefs, these are dereferenced and then unified (sometimes with copying) -}
        unify (HeapRef hr) (HeapRef gr)
          | BoundRef ha _ <- deref hr
          , BoundRef ga _ <- deref gr
          , ha == ga = uok
          | FreeRef ha <- deref hr
          , BoundRef ga _ <- deref gr = setHeap ha (HeapRef ga)
          | BoundRef ha _ <- deref hr
          , FreeRef ga <- deref gr = setHeap ga (HeapRef ha)
          | FreeRef ha <- deref hr
          , FreeRef ga <- deref gr = setHeap ha (HeapRef ga)
        unify (HeapRef hr') g =
          case deref hr' of
            FreeRef hr ->
              case g of
                atom@(Atom _) -> setHeap hr atom
                number@(Number _) -> setHeap hr number
                s@(Struct _) ->
                  withNewHeapStruct
                    hr
                    s
                    heap
                    (\nhs nheap ->
                       c
                         i
                           { cur =
                               cur
                                 {hed = map U 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 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
                        }
                  }
            BoundRef _ x -> unify x g
            _ -> ifail "dangling head ref"
        unify h (HeapRef gr') =
          case deref gr' of
            FreeRef gr ->
              case h of
                s@(Struct _) ->
                  withNewHeapStruct
                    gr
                    s
                    heap
                    (\ngs nheap ->
                       c
                         i
                           { cur =
                               cur
                                 {hed = hs, gol = map U ngs ++ gs, heap = nheap}
                           })
                x -> setHeap gr x
            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
                        }
                  }
            BoundRef _ x -> unify h x
            _ -> ifail "dangling goal ref"
        unify _ _ = backtrack
    {- Resolution -}
    go i@Interp { cur = cur@Cho { hed = hed
                                , hvar = hvar
                                , gol = gol
                                , gvar = gvar
                                , heap = heap
                                , stk = stk
                                , cut = cut
                                }
                , cho = chos
                }
      {- invoke a built-in (this gets replaced by NoGoal by default but the
       - builtin can actually do whatever it wants with the code) -}
      | [Invoke (Builtin bf)] <- hed =
        St.put i {cur = cur {hed = [NoGoal]}} >> bf
      {- top-level success -}
      | [NoGoal] <- hed
      , Just nchos <- tailcut gol chos cut
      , [] <- stk = do
        St.put i {cur = cur {hed = [], gol = []}, cho = nchos}
        finish $ 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 i =
      ifail $
      "code broken: impossible instruction combo hed=" ++
      show (hed . cur $ i) ++
      " gol=" ++ show (gol . cur $ i) ++ " stk=" ++ show (stk . cur $ i)
-}