{-# LANGUAGE RankNTypes #-}

module Interpreter where

{- pražský přehledný stroj -}
import Code
  ( Builtin(..)
  , Cho(..)
  , Code
  , Datum(..)
  , Instr(..)
  , InterpFn
  , emptyHeap
  , emptyScope
  )
import CodeLens
import Control.Monad (when)
import qualified Data.Map as M
import Env (PrlgEnv)
import Heap
import IR (Id(..), StrTable(..))
import Lens.Family2
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 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):_, _) -> advanceHead >> 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 = doCut >> advanceHead

cutGoal = doCut >> advance

openGoal :: IR.Id -> InterpFn
openGoal fn = do
  def <- defs `uses` (M.!? fn)
  case def of
    Just hs@(_:_) -> do
      advance
      cur . hvar .= emptyScope
      cur . unis .= arity fn
      cc <- use cur
      let (newcur:newcho) = [cc & 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)
  sretcut <- use (cur . retcut)
  cur . stk %= ((sgol, sgvar, scut, sretcut) :)
  cur . gol .= ngol
  cur . gvar .= ngvar
  cur . hed .= []
  cur . hvar .= emptyScope
  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
  continue

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

succeedHead :: InterpFn
succeedHead = do
  cur . hvar .= emptyScope
  cur . hed .= []
  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 (Atom _) VoidRef = uOK
unify VoidRef (Atom _) = uOK
unify (Atom a) (Atom b)
  | a == b = uOK
unify (Number _) VoidRef = uOK
unify VoidRef (Number _) = uOK
unify (Number a) (Number b)
  | a == b = uOK
unify (Struct a) VoidRef = do
  uNext
  cur . gol %= (replicate (arity a) (U VoidRef) ++)
  continue
unify VoidRef (Struct a) = do
  uNext
  cur . hed %= (replicate (arity a) (U VoidRef) ++)
  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' <- (cur . store) `uses` (M.!? lr)
  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
{- 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 -- BUG, structs!
          | 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)
-}