{-# LANGUAGE RankNTypes #-}

module Interpreter where

{- pražský přehledný stroj -}
import Code
  ( Builtin(..)
  , Cho(..)
  , Code
  , Datum(..)
  , Id(..)
  , Instr(..)
  , InterpFn
  , emptyHeap
  , emptyScope
  )
import CodeLens
import Control.Monad (when)
import qualified Data.Map as M
import Env (PrlgEnv)
import Heap
import IR (StrTable(..))
import Lens.Micro
import Lens.Micro.Mtl

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 = []
      , _hcut = []
      }
  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)
  {- tracing:
  import Control.Monad.Trans.Class (lift)
  import System.Console.Haskeline
  g <- use (cur . gol)
  cho <- use cho
  cut <- use (cur . cut)
  lift $ do
    outputStrLn $ "STEP (unis="++show u++")"
    outputStrLn $ "head = "++ show h
    outputStrLn $ "goal = "++ show g
    outputStrLn $ "cut = " ++ show cut
    outputStrLn $ "cho = " ++ show cho
  -}
  case (u, h) of
    (0, []) -> goalStep
    (0, _) -> headStep h
    (_, _)
      | u > 0 -> unifyStep h
    _ -> err "invalid interpreter state"

continue :: InterpFn
continue = pure Nothing

finish :: Bool -> InterpFn
finish = pure . Just . Right

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

{- toplevel choices -}
goalStep :: InterpFn
goalStep = do
  g <- use (cur . gol)
  case g of
    U (Struct s):gs -> openGoal s
    [Done] -> succeedGoal
    Cut:gs -> cutGoal
    Choices cs:gs -> pushChoices cs
    _ -> err "invalid goal code"

headStep :: [Instr] -> InterpFn
headStep h = do
  g <- use (cur . gol)
  case (h, g) of
    ([Done], _) -> succeedHead
    (Cut:_, _) -> cutHead
    (Invoke (Builtin bf):_, _) -> cur . hed .= [Done] >> bf
    (_, [Done]) -> tailCall
    (_, [Cut, Done]) -> tailCut
    (_, _) -> pushCall

unifyStep h = do
  g <- use (cur . gol)
  case (h, g) of
    (U hd:_, U gd:_) -> unify 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
      continue
    {- if there's no other choice available, answer no -}
    _ -> finish False

advance = do
  cur . gol %= tail
  continue

advanceHead = do
  cur . hed %= tail
  continue

{- resolution steps -}
doCut = use (cur . cut) >>= assign cho

retCut = do
  rc <- use (cur . retcut)
  when rc $ do
    doCut
    cur . retcut .= False

cutHead = do
  use (cur . hcut) >>= assign cho
  advanceHead

cutGoal = doCut >> advance

openGoal :: Id -> InterpFn
openGoal fn = do
  def <- (M.!? fn) <$> use defs
  case def of
    Just hs@(_:_) -> do
      advance
      cur . hvar .= emptyScope
      cur . unis .= arity fn
      cc <- use cur
      oldcho <- use cho
      let (newcur:newcho) = [cc & hcut .~ oldcho & hed .~ h | h <- hs]
      cur .= newcur
      cho %= (newcho ++)
      continue
    _ -> do
      StrTable _ _ itos <- use strtable
      err $ "no definition: '" ++ (itos M.! str fn) ++ "'/" ++ show (arity fn)

pushCall :: InterpFn
pushCall = do
  sgol <- use (cur . gol)
  sgvar <- use (cur . gvar)
  ngol <- use (cur . hed)
  ngvar <- use (cur . hvar)
  scut <- use (cur . cut)
  ncut <- use (cur . hcut)
  sretcut <- use (cur . retcut)
  cur . stk %= ((sgol, sgvar, scut, sretcut) :)
  cur . gol .= ngol
  cur . gvar .= ngvar
  cur . cut .= ncut
  cur . hed .= []
  cur . hvar .= emptyScope
  cur . hcut .= []
  cur . retcut .= False
  continue

tailCall :: InterpFn
tailCall = do
  ngol <- use (cur . hed)
  ngvar <- use (cur . hvar)
  cur . gol .= ngol
  cur . gvar .= ngvar
  cur . hed .= []
  cur . hvar .= emptyScope
  cur . hcut .= []
  continue

tailCut :: InterpFn
tailCut = do
  cur . retcut .= True
  advance
  tailCall

succeedHead :: InterpFn
succeedHead = do
  cur . hed .= []
  cur . hvar .= emptyScope
  cur . hcut .= []
  continue

succeedGoal :: InterpFn
succeedGoal = do
  retCut
  st <- use (cur . stk)
  case st of
    [] -> do
      cur . gol .= []
      finish True
    ((sgol, sgvar, scut, sretcut):st') -> do
      zoom cur $ do
        gol .= sgol
        gvar .= sgvar
        cut .= scut
        retcut .= sretcut
        stk .= st'
      continue

pushChoices :: [Code] -> InterpFn
pushChoices cs = do
  advance
  g <- use (cur . gol)
  let (ng:ogs) = [c ++ g | c <- cs]
  cc <- use cur
  cur . gol .= ng
  cho %= ([cc & gol .~ og | og <- ogs] ++)
  continue

{- unification -}
uNext = do
  advanceHead
  advance
  cur . unis -= 1

uOK :: InterpFn
uOK = uNext >> continue

unify :: Datum -> Datum -> InterpFn
unify VoidRef VoidRef = uOK
unify (C _) VoidRef = uOK
unify VoidRef (C _) = uOK
unify (C a) (C b)
  | a == b = uOK
  | a == b = uOK
unify (Struct a) VoidRef = do
  uNext
  cur . gol %= (replicate (arity a) (U VoidRef) ++)
  cur . unis += arity a
  continue
unify VoidRef (Struct a) = do
  uNext
  cur . hed %= (replicate (arity a) (U VoidRef) ++)
  cur . unis += arity a
  continue
unify (Struct a) (Struct b)
  | a == b = do
    cur . unis += arity a
    uOK
unify (LocalRef _) VoidRef = uOK
unify VoidRef (LocalRef _) = uOK
unify (LocalRef lr) g = do
  r <- findLocalRef hvar lr
  unify (HeapRef r) g
unify h (LocalRef lr) = do
  r <- findLocalRef gvar lr
  unify h (HeapRef r)
unify VoidRef (HeapRef _) = uOK
unify (HeapRef _) VoidRef = uOK
unify (HeapRef hr) (HeapRef gr) = do
  [h, g] <- traverse deref [hr, gr]
  case (h, g) of
    (BoundRef ha _, BoundRef ga _)
      | ha == ga -> uOK
    (BoundRef ha hv@(Struct Id {arity = arity}), BoundRef ga gv@(Struct _)) ->
      if hv /= gv
        then backtrack
        else do
          writeHeap ha (HeapRef ga) -- cycle unification trick thanks to Bart Demoen
          uNext
          cur . hed %= ([U . HeapRef $ ha + i | i <- [1 .. arity]] ++)
          cur . gol %= ([U . HeapRef $ ga + i | i <- [1 .. arity]] ++)
          cur . unis += arity
          continue
    (BoundRef _ hv, BoundRef _ gv)
      | hv == gv -> uOK
    (FreeRef ha, FreeRef ga) -> writeHeap ha (HeapRef ga) >> uOK
    (FreeRef ha, BoundRef ga _) -> writeHeap ha (HeapRef ga) >> uOK
    (BoundRef ha _, FreeRef ga) -> writeHeap ga (HeapRef ha) >> uOK
    _ -> backtrack
unify s@(Struct _) (HeapRef gr) = setStruct gr s gol
unify (HeapRef hr) s@(Struct _) = setStruct hr s hed
unify (Struct sa) (Struct sb)
  | sa == sb = cur . unis += arity sa >> uOK
unify h (HeapRef gr) = setSimple gr h
unify (HeapRef hr) g = setSimple hr g
unify _ _ = backtrack

{- unification reference-handling tools -}
findLocalRef :: Lens' Cho (M.Map Int Int) -> Int -> PrlgEnv Int
findLocalRef store lr = do
  a' <- (M.!? lr) <$> use (cur . store)
  case a' of
    Nothing -> do
      a <- newHeapVar
      cur . store %= M.insert lr a
      pure a
    Just a -> pure a

setStruct :: Int -> Datum -> Lens' Cho Code -> InterpFn
setStruct addr s@(Struct Id {arity = arity}) code = do
  x <- deref addr
  let cont nc = do
        uNext
        cur . unis += arity
        cur . code %= (map U nc ++)
        continue
  case x of
    FreeRef a -> putHeapStruct a s >>= cont
    BoundRef a s'@(Struct _)
      | s == s' -> cont [HeapRef (a + i) | i <- [1 .. arity]]
    _ -> backtrack

setSimple addr val = do
  x <- deref addr
  case x of
    FreeRef a -> writeHeap a val >> uOK
    BoundRef _ val'
      | val == val' -> uOK
    _ -> backtrack