Unison runtime prototyping (by Dan Doel)

rt.hs

--- Code by Dan Doel - https://github.com/unisonweb/unison/issues/1055#issuecomment-565753502
{-# language BangPatterns #-}

module POC (C(..), eval0, setup) where

import Control.Monad.Primitive

import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Unboxed.Mutable as UM

data Args
  = Arg1 !Int
  | Arg2 !Int !Int
  -- | ArgN {-# UNPACK #-} !(U.Vector Int)

data C
  = App !Int !Closure !Args
  | Lit !Int
  | Dec !Int
  | Add !Int !Int
  | Prn !Int
  | Let !C !C
  | Nop
  | Brn !Int !C !C
  | Ret !Int !Int
  | AddThn !Int !Int !C
  | LitThn !Int !C
  | DecThn !Int !C

data Closure
  = Ref !Int

data K = E | !C :< !K

type Stk = UM.MVector (PrimState IO) Int
type SP = Int

eval0 :: C -> IO ()
eval0 c = do
  stk <- GM.new 512
  eval stk 512 E c

resolve :: Closure -> C
resolve (Ref 0) = loop
resolve (Ref 1) = fib
resolve _ = error "Unknown reference"

eval :: Stk -> SP -> K -> C -> IO ()
eval !stk !sp !k  (Dec i) = do
  x <- GM.unsafeRead stk (sp + i)
  let sp' = sp-1
  GM.unsafeWrite stk sp' (x-1)
  jump stk sp' k
eval !stk !sp !k  (Add i j) = do
  x <- GM.unsafeRead stk (sp + i)
  y <- GM.unsafeRead stk (sp + j)
  let sp' = sp-1
  GM.unsafeWrite stk sp' (x+y)
  jump stk sp' k
eval !stk !sp !k  (Prn i) = do
  x <- GM.unsafeRead stk (sp + i)
  print x
  jump stk sp k
eval !stk !sp !k  (App lo f args) = do
  sp' <- moveArgs stk sp lo args
  eval stk sp' k (resolve f)
eval !stk !sp !k  (Lit n) = do
  let sp' = sp-1
  GM.unsafeWrite stk sp' n
  jump stk sp' k
eval !stk !sp !k  (Let c d) = do
  eval stk sp (d :< k) c
eval !stk !sp !k  (Brn i t f) = do
  b <- GM.unsafeRead stk (sp+i)
  if b == 0
    then eval stk sp k t
    else eval stk sp k f
eval !stk !sp !k  (Ret lo i) = do
  let sp' = sp+lo-1
  GM.unsafeRead stk (sp+i) >>= GM.unsafeWrite stk sp'
  jump stk sp' k
eval !stk !sp !k  (DecThn i d) = do
  x <- GM.unsafeRead stk (sp + i)
  let sp' = sp-1
  GM.unsafeWrite stk sp' (x-1)
  eval stk sp' k d
eval !stk !sp !k  (AddThn i j d) = do
  x <- GM.unsafeRead stk (sp + i)
  y <- GM.unsafeRead stk (sp + j)
  let sp' = sp-1
  GM.unsafeWrite stk sp' (x+y)
  eval stk sp' k d
eval !stk !sp !k  (LitThn n d) = do
  let sp' = sp-1
  GM.unsafeWrite stk sp' n
  eval stk sp' k d
eval !stk !sp !k Nop = jump stk sp k

jump :: Stk -> SP -> K -> IO ()
jump _   _  E = return ()
jump stk sp (d :< k) = eval stk sp k d
{-# inline jump #-}

moveArgs :: Stk -> SP -> Int -> Args -> IO Int
moveArgs !stk !sp !lo (Arg1 i) = do
  let sp' = sp + lo - 1
  GM.unsafeRead stk (sp + i) >>= GM.unsafeWrite stk sp'
  return sp'
moveArgs !stk !sp !lo (Arg2 i j) = do
  x <- GM.unsafeRead stk (sp + i)
  y <- GM.unsafeRead stk (sp + j)
  let sp' = sp + lo - 2
  GM.unsafeWrite stk sp' x
  GM.unsafeWrite stk (sp'+1) y
  return sp'
{-# inline moveArgs #-}

loop :: C
loop = Brn 0 fls tru
  where
  fls = Ret 2 1
  tru = Let (Add 0 1) $ Let (Dec 1) $ App 4 (Ref 0) (Arg2 0 1)

fib :: C
fib = Brn 0
        (Let (Lit 0) (Ret 2 0)) $
         Let (Dec 0) $
         Brn 0
           (Let (Lit 1) (Ret 3 0)) $
            Let (App 0 (Ref 1) (Arg1 0)) $
            Let (Dec 1) $
            Let (App 0 (Ref 1) (Arg1 0)) $
            Let (Add 0 2) $
            Ret 6 0

setup :: C
setup = Let (Lit 0) $ Let (Lit 100000000) $ Let (App 2 (Ref 0) (Arg2 0 1)) $ (Prn 0)
-- setup = Let (Lit 36) $ Let (App 1 (Ref 1) (Arg1 0)) $ Prn 0