Interp.hs

{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE TypeSynonymInstances  #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}

-- | Interpreter for L1
module Gibbon.L1.Interp ( interpProg, interp, applyPrim ) where

import           Data.ByteString.Builder ( toLazyByteString, string8)
import qualified Data.ByteString.Lazy.Char8 as B
import           Data.Char ( ord )
import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.State
import           Control.Monad.Writer
import qualified Data.List as L
import qualified Data.Map as M
import           System.Clock
import           System.IO.Unsafe
import           System.Random
import           Text.PrettyPrint.GenericPretty

import           Gibbon.Common
import           Gibbon.L1.Syntax as L1


interpChatter :: Int
interpChatter = 7

------------------------------------------------------------

instance InterpExt () Exp1 (E1Ext () Ty1) Var where
  gInterpExt rc valenv ddefs fundefs ex =
      case ex of
          BenchE fn locs args _b -> interp rc valenv ddefs fundefs (AppE fn locs args)
          AddFixed{}   -> error "L1.Interp: AddFixed not handled."
          StartOfPkdCursor{} -> error "L1.Interp: StartOfPkdCursor not handled."

instance Interp () Exp1 Var where
  gInterpExp = interp

instance InterpProg () Exp1 Var where
  gInterpProg _s = interpProg

-- | Interpret a program, including printing timings to the screen.
--   The returned bytestring contains that printed timing info.
interpProg :: Interp () e Var => RunConfig -> Prog Var e -> IO ((), Value e, B.ByteString)
interpProg rc Prog{ddefs,fundefs,mainExp} =
  case mainExp of
    -- Print nothing, return "void"
    Nothing -> return ((), VProd [], B.empty)
    Just (e,_) -> do
      let fenv = M.fromList [ (funName f , f) | f <- M.elems fundefs]
      (v,logs,_) <- runInterpM (gInterpExp rc M.empty ddefs fenv e) ()
      pure ((), v, toLazyByteString logs)

interp :: forall e l d s.
          (Show l, Ord l, NFData l, Out l,
           Show d, NFData d, Out d, Ord d,
           Ord (e l d), NFData (e l d),
           InterpExt s (PreExp e l d) (e l d) Var)
       => RunConfig
       -> ValEnv Var (PreExp e l d)
       -> DDefs (TyOf (PreExp e l d))
       -> FunDefs Var (PreExp e l d)
       -> (PreExp e l d)
       -> InterpM s (PreExp e l d) (Value (PreExp e l d))
interp rc valenv ddefs fenv = go valenv
  where
    {-# NOINLINE goWrapper #-}
    goWrapper env !_ix ex = go env ex

    go :: ValEnv Var (PreExp e l d) -> (PreExp e l d) -> InterpM s (PreExp e l d) (Value (PreExp e l d))
    go env x0 = do
        case x0 of
          Ext ext -> do
              gInterpExt rc env ddefs fenv ext

          LitE c    -> return $ VInt c
          CharE c   -> return $ VChar c
          FloatE c  -> return $ VFloat c
          LitSymE s -> return $ VSym (fromVar s)
          VarE v    -> do
              return $ env # v

          -- -- Don't sort for now
          -- PrimAppE (VSortP{}) [ls,VarE fp] -> do
          --   (VList vals) <- go env ls
          --   applySortP env vals fp

          PrimAppE p ls -> do args <- mapM (go env) ls
                              applyPrim rc p args
          ProjE ix ex   -> do VProd ls <- go env ex
                              return $ ls !!! ix

          -- N.B. this AppE is shared by the interpreters for L0 and L1
          AppE f _ ls -> do
            ls' <- mapM (go env) ls
            -- Look in the local environment first
            case M.lookup f env of
              Nothing ->
                case M.lookup f fenv of
                  Just fn -> do
                      let env' = M.union (M.fromList (zip (funArgs fn) ls')) env
                      go env' (funBody fn)
                  Nothing -> error $ "L1.Interp: unbound function in application: "++ndoc x0
              Just fn@(VLam args bod closed_env) ->
                if length args /= length ls'
                then error $ "L0.Interp: unexpected arguments in application: "++ ndoc ls' ++ " in " ++ ndoc fn
                else do
                  let env' = M.fromList (zip args ls') `M.union` closed_env `M.union` env
                  go env' bod
              Just oth -> error $ "L0.Interp: expected a lambda in application, got: "++ ndoc oth ++ " in " ++ ndoc oth


          CaseE _ [] -> error$ "L1.Interp: CaseE with empty alternatives list: "++ndoc x0

          CaseE x1 alts@((_sometag,_,_):_) -> do
                 v <- go env x1
                 case v of
                   VPacked k ls2 -> do
                       let vs = L.map fst prs
                           (_,prs,rhs) = lookup3 k alts
                           env' = M.union (M.fromList (zip vs ls2)) env
                       go env' rhs
                   _ -> error$ "L1.Interp: type error, expected data constructor, got: "++ndoc v++
                               "\nWhen evaluating scrutinee of case expression: "++ndoc x1

          LetE (v,_,_ty,rhs) bod -> do
            rhs' <- go env rhs
            let env' = M.insert v rhs' env
                env'' = case rhs of
                          (PrimAppE (InplaceVUpdateP _) [VarE x,_,_]) ->
                            M.insert x rhs' env'
                          _ -> env'
            go env'' bod

          MkProdE ls -> VProd <$> mapM (go env) ls
          -- TODO: Should check this against the ddefs.
          DataConE _ k ls -> do
              args <- mapM (go env) ls
              return $ VPacked k args

          TimeIt bod _ isIter -> do
              let iters = if isIter then rcIters rc else 1
              !_ <- return $! force env
              st <- liftIO $ getTime clk
              val <- foldM (\ _ i -> goWrapper env i bod)
                            (error "Internal error: this should be unused.")
                         [1..iters]
              en <- liftIO $ getTime clk
              let tm = fromIntegral (toNanoSecs $ diffTimeSpec en st)
                        / 10e9 :: Double
              if isIter
               then do tell$ string8 $ "ITERS: "++show iters       ++"\n"
                       tell$ string8 $ "SIZE: " ++show (rcSize rc) ++"\n"
                       tell$ string8 $ "BATCHTIME: "++show tm      ++"\n"
               else tell$ string8 $ "SELFTIMED: "++show tm ++"\n"
              return $! val

          SpawnE f locs args -> go env (AppE f locs args)
          SyncE -> pure $ VInt (-1)

          WithArenaE v e -> do
              let env' = M.insert v (VInt 0) env
              go env' e

          IfE a b c -> do v <- go env a
                          case v of
                           VBool flg -> if flg
                                        then go env b
                                        else go env c
                           oth -> error$ "interp: expected bool, got: "++show oth

          MapE _ _bod    -> error "L1.Interp: finish MapE"
          FoldE _ _ _bod -> error "L1.Interp: finish FoldE"

    -- _applySortP :: ValEnv (PreExp e l d) -> [(Value (PreExp e l d))] -> Var -> WriterT InterpLog IO (Value (PreExp e l d))
    -- _applySortP env ls f = do
    --   let fn  = case M.lookup f fenv of
    --               Just fun -> fun
    --               Nothing -> error $ "L1.Interp: unbound function given to vsort: "++ndoc f
    --       ls' = sortBy
    --             (\a b ->
    --                  let env' = M.union (M.fromList (zip (funArgs fn) [a,b])) env
    --                      (i,_) =
    --                        unsafePerformIO $ (runWriterT (go env' (funBody fn)))
    --                      VInt j = i
    --                  in compare j 0)
    --             ls
    --   pure (VList ls')


applyPrim :: (Show ty, Ord l, Out l, Show l, Show d, Out d, Ord d,  Ord (e l d), Out (e l d), Show (e l d))
          => RunConfig -> Prim ty -> [(Value (PreExp e l d))] -> InterpM s (PreExp e l d) (Value (PreExp e l d))
applyPrim rc p args =
 case (p,args) of
   (MkTrue,[])             -> pure $ VBool True
   (MkFalse,[])            -> pure $ VBool False
   -- FIXME: randomIO does not guarentee unique numbers every time.
   (Gensym, [])            -> pure $ VSym $ "gensym_" ++ (show $ (unsafePerformIO randomIO :: Int) `mod` 1000)
   (AddP,[VInt x, VInt y]) -> pure $ VInt (x+y)
   (SubP,[VInt x, VInt y]) -> pure $ VInt (x-y)
   (MulP,[VInt x, VInt y]) -> pure $ VInt (x*y)
   (DivP,[VInt x, VInt y]) -> pure $ VInt (x `quot` y)
   (ModP,[VInt x, VInt y]) -> pure $ VInt (x `rem` y)
   (ExpP,[VInt x, VInt y]) -> pure $ VInt (x ^ y)
   (FAddP,[VFloat x, VFloat y]) -> pure $ VFloat (x+y)
   (FSubP,[VFloat x, VFloat y]) -> pure $ VFloat (x-y)
   (FMulP,[VFloat x, VFloat y]) -> pure $ VFloat (x*y)
   (FDivP,[VFloat x, VFloat y]) -> pure $ VFloat (x / y)
   (FExpP,[VFloat x, VFloat y]) -> pure $ VFloat (x ** y)
   -- Constrained to the value of RAND_MAX (in C) on my laptop: 2147483647 (2^31 − 1)
   (RandP,[]) -> do
       i <- liftIO $ randomIO
       pure $ VInt $ i `mod` 2147483647
   (FRandP,[]) -> do
       i <- liftIO $ randomIO
       pure $ VFloat i
   (IntToFloatP,[VInt x]) -> pure $ VFloat (fromIntegral x)
   (FloatToIntP,[VFloat x]) -> pure $ VInt (round x)
   (FSqrtP,[VFloat x]) -> pure $ VFloat (sqrt x)
   (EqSymP,[VSym x, VSym y]) -> pure $ VBool (x==y)
   (EqBenchProgP _str,[]) -> pure $ VBool False
   (EqIntP,[VInt x, VInt y]) -> pure $ VBool (x==y)
   (EqFloatP,[VFloat x, VFloat y]) -> pure $ VBool (x==y)
   (EqCharP ,[VChar x , VChar y])  -> pure $ VBool (x==y)
   (LtP,[VInt x, VInt y]) -> pure $ VBool (x < y)
   (GtP,[VInt x, VInt y]) -> pure $ VBool (x > y)
   (LtEqP,[VInt x, VInt y]) -> pure $ VBool (x <= y)
   (GtEqP,[VInt x, VInt y]) -> pure $ VBool (x >= y)
   (FLtP,[VFloat x, VFloat y]) -> pure $ VBool (x < y)
   (FGtP,[VFloat x, VFloat y]) -> pure $ VBool (x > y)
   (FLtEqP,[VFloat x, VFloat y]) -> pure $ VBool (x <= y)
   (FGtEqP,[VFloat x, VFloat y]) -> pure $ VBool (x >= y)
   (AndP, [VBool x, VBool y]) -> pure $ VBool (x && y)
   (OrP, [VBool x, VBool y])  -> pure $ VBool (x || y)
   ((DictInsertP _ty),[_, VDict mp, key, val]) -> pure $ VDict (M.insert key val mp)
   ((DictLookupP _),[VDict mp, key])        -> pure $ mp # key
   ((DictHasKeyP _),[VDict mp, key])        -> pure $ VBool (M.member key mp)
   ((DictEmptyP _),[_])                     -> pure $ VDict M.empty
   ((ErrorP msg _ty),[]) -> pure $ error msg
   (SizeParam,[]) -> pure $ VInt (rcSize rc)
   (IsBig,[_one,_two]) -> pure $ VBool False
   (ReadPackedFile file _ _ ty,[]) ->
       error $ "L1.Interp: unfinished, need to read a packed file: "++show (file,ty)
   (ReadArrayFile{},[]) -> do
       pure (VList [])
   (VAllocP _,_n) -> pure (VList [])
   (VFreeP _,_n) -> pure (VProd [])
   (VFree2P _,_n) -> pure (VProd [])
   (VLengthP _,[(VList ls)]) -> pure $ VInt (length ls)
   (VNthP _,[(VList ls), VInt n]) -> pure $ ls !!! n
   (InplaceVUpdateP _,[(VList ls), VInt i, v]) -> do
       let ls' = if length ls <= i
                 then
                     let need = (i+1) - (length ls)
                     in VList $ ls ++ (replicate need v)
                 else VList (replaceNth i v ls)
       pure ls'
   (VSliceP _,[VInt from, VInt len, (VList ls)]) -> do
       pure $ VList (L.take len (L.drop from ls))
   (VConcatP _, [VList ls]) -> do
       let concatd = foldl (\acc v ->
                                case v of
                                  VList l -> acc ++ l
                                  _ -> error $ "VConcatP: " ++ show v)
                           []
                           ls
       pure $ VList concatd
   (VMergeP _,[VList ls1,VList ls2]) -> do
       pure $ VList (ls1 ++ ls2)
   (GetNumProcessors, []) -> pure $ VInt 1
   -- Don't sort for now.
   (VSortP _, [ls, _fn]) -> do
       pure ls
   (InplaceVSortP _, [ls, _fn]) -> do
       pure ls
   (PrintInt, [VInt n]) -> do
       tell $ string8 (show n)
       pure $ VProd []
   (PrintFloat, [VFloat n]) -> do
       tell $ string8 (show n)
       pure $ VProd []
   (PrintBool, [VBool n]) -> do
       tell $ string8 (show n)
       pure $ VProd []
   (PrintSym, [VSym s]) -> do
       tell $ string8 s
       let v' = case s of
                  "NEWLINE" -> "\n"
                  "SPACE"   -> " "
                  _oth -> s
       liftIO $ putStr v'
       pure $ VProd []
   (PrintSym, [VInt n]) -> do
       tell $ string8 (show n)
       pure $ VProd []
   oth -> error $ "unhandled prim or wrong number of arguments: "++show oth

  where
     replaceNth :: Int -> a -> [a] -> [a]
     replaceNth _ _ [] = []
     replaceNth n newVal (x:xs)
       | n == 0 = newVal:xs
       | otherwise = x:replaceNth (n-1) newVal xs


clk :: Clock
clk = Monotonic

strToInt :: String -> Int
strToInt = product . L.map ord