a new impl. to fix some shortcomings from how we represent thunks

NewEval.hs

{-# LANGUAGE RecursiveDo               #-}
{-# LANGUAGE LambdaCase                #-}
{-# LANGUAGE DeriveGeneric             #-}
{-# LANGUAGE DuplicateRecordFields     #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TemplateHaskell           #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE UndecidableInstances      #-}
{-# LANGUAGE QuasiQuotes               #-}

module Language.Jsonnet.NewEval where

import Language.Jsonnet.Core
import Data.ByteString (ByteString)
import Language.Jsonnet.Common
import Data.Text (Text)
import Data.HashMap.Lazy (HashMap)
import qualified Data.HashMap.Lazy as H
import Data.Map.Lazy (Map)
import qualified Data.Map.Lazy as M
import Data.Vector (Vector, (!?))
import qualified Data.Vector as V
import Data.Scientific (Scientific, fromFloatDigits, toRealFloat)
import Data.Text (Text)
import qualified Data.Text as T
import Data.Text.Encoding (decodeUtf8, encodeUtf8)
import Unbound.Generics.LocallyNameless
import Control.Applicative
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Except
import Data.IntMap.Lazy (IntMap, (!))
import qualified Data.IntMap.Lazy as IntMap
import GHC.Generics
import Data.Maybe
import Control.Lens (makeLenses, (%=), (<+=), use)
import qualified Data.Aeson as JSON
import Unbound.Generics.LocallyNameless.Internal.Fold (foldMapOf, toListOf)
import Data.Monoid (Any(..))
import Data.Typeable (Typeable)
import Language.Jsonnet.TH.QQ
import Language.Jsonnet.Desugar
import Language.Jsonnet.Error
import Language.Jsonnet.Pretty ()
import Text.PrettyPrint.ANSI.Leijen hiding (encloseSep, (<$>))
import Data.Scientific (Scientific, fromFloatDigits, toRealFloat)
import Data.Bifunctor (second)
import Data.List (find)
import Language.Jsonnet.Annotate

type Loc = Int

type Error = String

type Eval a = ExceptT EvalError (ReaderT Env (StateT Heap (FreshMT IO))) a

data Value
  = VNull
  | VBool !Bool
  | VNum !Scientific
  | VStr !Text
  | VObj !Object
  | VArr !(Vector Value)
  | VThunk !Core !Env
  | VIndir !Loc
  | VFun !(Value -> Eval Value)
  | VPrim !Prim
  | VClos !Fun !Env

type Object = HashMap Text (Hideable Value)

--data Object = Object
--  { _self :: Loc,
--    _super :: Loc,
--    _dict :: HashMap Text (Hideable Value),
--  }

data Cell = Cell { cellVal :: Value, cellIsWHNF :: Bool }
  deriving (Generic)

mkCell :: Value -> Cell
mkCell v = Cell v False

type Env = Map (Name Core) Value

data Heap = Heap
  { _memory :: IntMap Cell,
    _nextLoc :: Loc
  }
  deriving (Generic)

makeLenses ''Heap

emptyHeap :: Heap
emptyHeap = Heap IntMap.empty 0

class HasValue a where
  inj :: a -> Value
  proj :: Value -> Eval a

instance HasValue Value where
  inj = id
  proj = pure

instance HasValue Bool where
  proj (VBool n) = pure n
  proj v = throwTypeMismatch "bool" v
  inj = VBool

instance HasValue Text where
  proj (VStr s) = pure s
  proj v = throwTypeMismatch "string" v
  inj = VStr

instance {-# OVERLAPPING #-} HasValue [Char] where
  proj (VStr s) = pure $ T.unpack s
  proj v = throwTypeMismatch "string" v
  inj = VStr . T.pack

instance HasValue ByteString where
  proj (VStr s) = pure (encodeUtf8 s)
  proj v = throwTypeMismatch "string" v
  inj = VStr . decodeUtf8

instance HasValue Scientific where
  proj (VNum n) = pure n
  proj v = throwTypeMismatch "number" v
  inj = VNum

instance HasValue Double where
  proj (VNum n) = pure (toRealFloat n)
  proj v = throwTypeMismatch "number" v
  inj = VNum . fromFloatDigits

instance {-# OVERLAPS #-} Integral a => HasValue a where
  proj (VNum n) = pure (round n)
  proj v = throwTypeMismatch "number" v
  inj = VNum . fromIntegral

instance HasValue a => HasValue (Maybe a) where
  proj VNull = pure Nothing
  proj a = Just <$> proj a
  inj Nothing = VNull
  inj (Just a) = inj a


instance HasValue a => HasValue (Vector a) where
  inj as = VArr (inj <$> as)
  proj (VArr as) = mapM (whnfV >=> proj) as

instance {-# OVERLAPPABLE #-} HasValue a => HasValue [a] where
  inj = inj . V.fromList
  proj = fmap V.toList . proj

instance {-# OVERLAPS #-} (HasValue a, HasValue b) => HasValue (a -> b) where
  inj f = VFun $ whnfV >=> fmap (inj . f) . proj

instance {-# OVERLAPS #-} (HasValue a, HasValue b, HasValue c) => HasValue (a -> b -> c) where
  inj f = inj $ \x -> inj (f x)

instance {-# OVERLAPS #-} (HasValue a, HasValue b) => HasValue (a -> Eval b) where
  inj f = VFun $ whnfV >=> proj >=> (fmap inj . f)
  proj (VFun f) = pure $ proj <=< f . inj

instance {-# OVERLAPS #-} (HasValue a, HasValue b, HasValue c) => HasValue (a -> b -> Eval c) where
  inj f = inj $ \x -> inj (f x)

inj'' ::
  (HasValue a, HasValue b, HasValue c) =>
  (a -> b -> c) ->
  Value ->
  Value ->
  Eval Value
inj'' f v1 v2 = inj <$> liftA2 f (proj v1) (proj v2)

-- try to go throw an example that showcases why this approach is good

rnf :: Core -> Eval Value
rnf c = whnf c >>= rnfV

rnfV :: Value -> Eval Value
rnfV (VObj vs) = VObj <$> mapM (mapM rnfV) vs
rnfV (VArr vs) = VArr <$> mapM rnfV vs
rnfV v@(VThunk {}) = whnfV v >>= rnfV
rnfV v@(VIndir {}) = whnfV v >>= rnfV
rnfV v = pure v

whnfV :: Value -> Eval Value
whnfV (VThunk c e) = withEnv e (whnf c) >>= whnfV
whnfV (VIndir loc) = whnfIndir loc
whnfV v = pure v

whnf :: Core -> Eval Value
whnf (CVar n) = lookupVar n
whnf (CLoc _ c) = whnf c
whnf (CLit (String s)) = pure (VStr s)
whnf (CLit (Number n)) = pure (VNum n)
whnf (CObj dict) = whnfObj dict
whnf (CArr cs) = VArr . V.fromList <$> mapM mkValue cs
whnf (CLet bnd) = whnfLetrec bnd
whnf (CPrim p) = pure (VPrim p)
whnf (CApp e es) = whnfApp e es
whnf (CFun f) = VClos f <$> ask

whnfArgs :: Args Core -> Eval [Arg Value]
whnfArgs = \case
  as@(Args _ Strict) -> args <$> mapM whnf as
  as@(Args _ Lazy)   -> args <$> mapM mkValue as

whnfApp :: Core -> Args Core -> Eval Value
whnfApp e es = do
  vs <- whnfArgs es
  whnf e >>= \case
    VClos f env -> whnfClos env f vs
    VPrim op -> whnfPrim op vs
    v -> throwTypeMismatch "primitive" v

whnfClos :: Env -> Fun -> [Arg Value] -> Eval Value
whnfClos rho (Fun f) args = do
  (bnds, e) <- unbind f
  (rs, ps, ns) <- splitArgs args (second unembed <$> unrec bnds)
  withEnv rho $
    extendEnv (M.fromList ps) $
      extendEnv (M.fromList ns) $
        appDefaults rs e

-- all parameter names are bound in default values
appDefaults rs e = mdo
  bnds <-
    M.fromList <$> mapM
      ( \(n, e) -> do
          th <- extendEnv bnds (mkValue e)
          pure (n, th)
      )
      rs
  extendEnv bnds (whnf e)

-- returns a triple with unapplied binders, positional and named
splitArgs args bnds = do
  named <- getNamed
  pos <- getPos
  unapp <- getUnapp named
  pure (unapp, pos, named)
  where
    (bnds1, bnds2) = splitAt (length ps) bnds
    (ps, ns) = split args
    pNames = fst (unzip bnds)

    getPos = do
      if length ps > length bnds
        then throwError $ TooManyArgs (length bnds)
        else pure $ zip (fst $ unzip bnds1) ps

    -- checks the provided named arguments exist
    getNamed = traverse f ns
      where
        f (a, b) = case g a of
          Nothing -> throwError $ BadParam (pretty a)
          Just n -> pure (n, b)
        g a = find ((a ==) . name2String) pNames

    getUnapp named =
      pure $ filter ((`notElem` ns) . fst) bnds2
        where
          ns = fst (unzip named)

    split [] = ([], [])
    split (Pos p : xs) =
      let (ys, zs) = split xs in (p : ys, zs)
    split (Named n v : xs) =
      let (ys, zs) = split xs in (ys, (n, v) : zs)

whnfPrim :: Prim -> [Arg Value] -> Eval Value
whnfPrim (BinOp op) [Pos e1, Pos e2] =
  liftA2 (,) (whnfV e1) (whnfV e2) >>= uncurry (whnfBinOp op)
whnfPrim Cond [Pos c, Pos t, Pos e] = whnfCond c t e

whnfBinOp :: BinOp -> Value -> Value -> Eval Value
whnfBinOp Lookup e1 e2 = whnfLookup e1 e2
whnfBinOp Add (VObj x) (VObj y) = x `mergeWith` y
whnfBinOp Add n1 n2 = evalBin ((+) @Double) n1 n2
whnfBinOp Sub n1 n2 = evalBin ((-) @Double) n1 n2
whnfBinOp Eq n1 n2 = evalBin ((==) @Double) n1 n2
whnfBinOp Le e1 e2 = evalBin ((<=) @Double) e1 e2
whnfBinOp op _ _ = throwError (InvalidKey (pretty $ show op))

whnfCond :: Value -> Value -> Value -> Eval Value
whnfCond c e1 e2 = do
  c' <- (whnfV >=> proj) c
  if c'
    then whnfV e1
    else whnfV e2

evalBin ::
  (HasValue a, HasValue b, HasValue c) =>
  (a -> b -> c) ->
  Value ->
  Value ->
  Eval Value
evalBin = inj''

whnfLookup :: Value -> Value -> Eval Value
whnfLookup (VObj o) (VStr s) =
  liftMaybe (NoSuchKey (pretty s)) (H.lookup s o) >>= \(Hideable v _) -> pure v

throwTypeMismatch :: Text -> Value -> Eval a
throwTypeMismatch e f = throwError =<< TypeMismatch e <$> valueType f

lookupVar :: Name Core -> Eval Value
lookupVar n = do
  rho <- ask
  v <- liftMaybe (VarNotFound (pretty n)) (M.lookup n rho)
  whnfV v

--mkIndir :: Core -> Eval Value
--mkIndir c@(CLit _) = whnf c
--mkIndir c = VIndir <$> (allocate =<< (VThunk c <$> ask))

mkIndirV :: Value -> Eval Value
mkIndirV v = VIndir <$> allocate v

mkThunk :: Core -> Eval Value
mkThunk c = VThunk c <$> ask

mkValue :: Core -> Eval Value
mkValue c@(CLit _) = whnf c
mkValue c = mkThunk c >>= mkIndirV

-- this comes from Disco!
whnfIndir :: Loc -> Eval Value
whnfIndir loc = do
  m <- use memory                   -- Get the memory map
  let c = m ! loc                   -- Look up the given location and reduce it to WHNF
  case c of
    Cell v True  -> return v        -- Already evaluated, just return it
    Cell v False -> do
      v' <- whnfV v                               -- Needs to be reduced
      memory %= IntMap.insert loc (Cell v' True)  -- Update memory with the reduced value
      return v'                                   -- Finally, return the value.

whnfLetrec :: Let -> Eval Value
whnfLetrec (Let bnd) = mdo
  (r, e1) <- unbind bnd
  bnds <-
    M.fromList <$> mapM
      ( \(n, Embed e) -> do
          v <- extendEnv bnds (mkValue e)
          pure (n, v)
      )
      (unrec r)
  extendEnv bnds (mkValue e1)

whnfObj :: [KeyValue Core] -> Eval Value
whnfObj xs = mdo
  self <-
    mkIndirV =<< VObj . H.fromList . catMaybes
      <$> mapM
        ( \(KeyValue key val) -> do
            let bnd = M.fromList [(s2n "self", self)]
            k <- whnf key
            v <- whnfObjValue bnd `traverse` val
            case k of
              VStr k -> pure $ Just (k, v)
              _ -> pure Nothing
        )
        xs
  whnfV self

-- | Right-biased union of two objects, i.e. '{x : 1} + {x : 2} == {x : 2}'
--   with OO-like `self` and `super` support via value recursion (knot-tying)
mergeWith :: Object -> Object -> Eval Value
mergeWith xs ys = mdo
  zs' <- mkIndirV $ VObj (H.unionWith f xs' ys')
  xs' <- pure $ self zs' <$> xs
  ys' <- do
    xs'' <- mkIndirV (VObj xs')
    mapM (super xs'') (self zs' <$> ys)
  pure zs'
  where
    f a b
      | hidden a && visible b = a
      | otherwise = b
    self xs = fmap $ \case
      VThunk c env ->
        let env' = M.insert "self" xs env
        in VThunk c env'
      v -> v
    super xs = mapM $ \case
      VThunk c env ->
        let env' = M.insert "super" xs env
        in
          if hasSuper c
            then mkIndirV (VThunk c env') -- memoized the result
            else pure (VThunk c env')
      v -> pure v

whnfObjValue :: Env -> Core -> Eval Value
whnfObjValue _ c@(CLit _) = whnf c
whnfObjValue self e
  | lb e = extendEnv self (mkThunk e)
  | otherwise = mkValue e
  where
    lb e = hasSelf e || hasSuper e

hasSelf, hasSuper :: Core -> Bool
hasSelf = isFreeIn (s2n "self")
hasSuper = isFreeIn (s2n "super")

isFreeIn :: Alpha b => Name Core -> b -> Bool
isFreeIn = elementOf fv
  where
    elementOf l = anyOf l . (==)
    anyOf l f = getAny . foldMapOf l (Any . f)

extendEnv :: Env -> Eval a -> Eval a
extendEnv = local . M.union

withEnv :: Env -> Eval a -> Eval a
withEnv = local . const

allocate :: Value -> Eval Loc
allocate v = do
  loc <- nextLoc <+= 1
  -- io $ putStrLn $ "allocating " ++ show v ++ " at location " ++ show loc
  memory %= IntMap.insert loc (mkCell v)
  return loc

visibleKeys :: Object -> HashMap Text Value
visibleKeys = H.mapMaybe f
  where
    f vv@(Hideable v _)
      | not (hidden vv) = Just v
      | otherwise = Nothing

valueType :: Value -> Eval Text
valueType (VStr _) = pure "string"
valueType (VObj _) = pure "object"
valueType (VArr _) = pure "array"
valueType (VClos {}) = pure "function"
valueType (VFun _) = pure "function"
valueType (VPrim _) = pure "function"
valueType v@(VThunk {}) = whnfV v >>= valueType
valueType v@(VIndir {}) = whnfV v >>= valueType

manifest :: Value -> Eval JSON.Value
manifest (VStr s) = pure (JSON.String s)
manifest (VNum s) = pure (JSON.Number s)
manifest (VArr vs) = JSON.Array <$> mapM manifest vs
manifest (VObj vs) = JSON.Object <$> mapM manifest (visibleKeys vs)


foo = [jsonnet|
  local fibnext = {
    a: super.a + super.b,
    b: super.a,
  };
  local fib(n) =
    if n == 0 then
      { a: 1, b: 1 }
    else
      fib(n - 1) + fibnext;

  fib(25)
|]

bar = desugar (annMap (const ()) foo)

result = rnf bar >>= manifest

run e = runFreshMT (evalStateT (runReaderT (runExceptT e) M.empty) emptyHeap)

liftMaybe :: EvalError -> Maybe a -> Eval a
liftMaybe e =
  \case
    Nothing -> throwError e
    Just a -> pure a