<commit_before>{-# LANGUAGE GeneralizedNewtypeDeriving, ScopedTypeVariables #-}
module L0.Interpreter
  ( runProg
  , runProgIO
  , InterpreterError(..) )
where

import Control.Applicative
import Control.Monad.Reader
import Control.Monad.Trans.Either

import Data.Array
import Data.Bits
import Data.List
import Data.Loc
import qualified Data.Map as M

import System.IO (hFlush, stdout)

import L0.AbSyn
import L0.Parser

data InterpreterError = MissingMainFunction
                      | IndexOutOfBounds Loc Int Int
                      -- ^ First @Int@ is array size, second is attempted index.
                      | NegativeIota Loc Int
                      | NegativeReplicate Loc Int
                      | ReadError Loc Type String
                      | InvalidArrayShape Loc [Int] [Int]
                      -- ^ First @Int@ is old shape, second is attempted new shape.
                      | TypeError Loc String

instance Show InterpreterError where
  show MissingMainFunction =
    "No main function defined."
  show (IndexOutOfBounds pos arrsz i) =
    "Array index " ++ show i ++ " out of bounds of array size " ++ show arrsz ++ " at " ++ locStr pos ++ "."
  show (NegativeIota pos n) =
    "Argument " ++ show n ++ " to iota at " ++ locStr pos ++ " is negative."
  show (NegativeReplicate pos n) =
    "Argument " ++ show n ++ " to replicate at " ++ locStr pos ++ " is negative."
  show (TypeError pos s) =
    "Type error at " ++ locStr pos ++ " in " ++ s ++ " during interpretation.  This implies a bug in the type checker."
  show (ReadError pos t s) =
    "Read error while trying to read " ++ ppType t ++ " at " ++ locStr pos ++ ".  Input line was: " ++ s
  show (InvalidArrayShape pos shape newshape) =
    "Invalid array reshaping at " ++ locStr pos ++ ", from " ++ show shape ++ " to " ++ show newshape


data L0Env m = L0Env { envVtable  :: M.Map String Value
                     , envFtable  :: M.Map String ([Value] -> L0M m Value)
                     , envWriteOp :: String -> L0M m ()
                     , envReadOp  :: L0M m String }

newtype L0M m a = L0M (ReaderT (L0Env m) (EitherT InterpreterError m) a)
  deriving (MonadReader (L0Env m), Monad, Applicative, Functor)

runL0M :: L0M m a -> L0Env m -> m (Either InterpreterError a)
runL0M (L0M m) env = runEitherT (runReaderT m env)

bad :: Monad m => InterpreterError -> L0M m a
bad = L0M . lift . left

bindVar :: L0Env m -> (String, Value) -> L0Env m
bindVar env (name,val) =
  env { envVtable = M.insert name val $ envVtable env }

bindVars :: L0Env m -> [(String, Value)] -> L0Env m
bindVars = foldl bindVar

binding :: Monad m => [(String, Value)] -> L0M m a -> L0M m a
binding bnds = local (`bindVars` bnds)

lookupVar :: Monad m => String -> L0M m Value
lookupVar fname = do val <- asks $ M.lookup fname . envVtable
                     case val of Just val' -> return val'
                                 Nothing   -> bad $ TypeError noLoc "lookupVar"

lookupFun :: Monad m => String -> L0M m ([Value] -> L0M m Value)
lookupFun fname = do fun <- asks $ M.lookup fname . envFtable
                     case fun of Just fun' -> return fun'
                                 Nothing   -> bad $ TypeError noLoc "lookupFun"

arrToList :: Monad m => Value -> L0M m [Value]
arrToList (ArrayVal l _ _) = return $ elems l
arrToList v = bad $ TypeError (locOf v) "arrToList"

tupToList :: Monad m => Value -> L0M m [Value]
tupToList (TupVal l _) = return l
tupToList v = bad $ TypeError (locOf v) "tupToList"

runProgIO :: Prog Type -> IO (Either InterpreterError Value)
runProgIO = runProg putStr (hFlush stdout >> getLine)

runProg :: (Applicative m, Monad m) => (String -> m ()) -> m String
        -> Prog Type -> m (Either InterpreterError Value)
runProg wop rop prog = do
  let ftable = foldl expand builtins prog
      l0env = L0Env { envVtable = M.empty
                    , envFtable = ftable
                    , envWriteOp = L0M . lift . lift . wop
                    , envReadOp = L0M $ lift $ lift rop }
      runmain = case M.lookup "main" ftable of
                  Nothing -> bad MissingMainFunction
                  Just mainfun -> mainfun []
  runL0M runmain l0env
  where
    -- We assume that the program already passed the type checker, so
    -- we don't check for duplicate definitions or anything.
    expand ftable (name,_,params,body,_) =
        let fun args = binding (zip (map fst params) args) $ evalExp body
        in M.insert name fun ftable
    builtins = M.fromList [("toReal", builtin "toReal")
                          ,("trunc", builtin "trunc")
                          ,("sqrt", builtin "sqrt")
                          ,("log", builtin "log")
                          ,("exp", builtin "exp")
                          ,("op not", builtin "op not")
                          ,("op ~", builtin "op ~")]
    builtin "toReal" [IntVal x pos] = return $ RealVal (fromIntegral x) pos
    builtin "trunc" [RealVal x pos] = return $ IntVal (truncate x) pos
    builtin "sqrt" [RealVal x pos] = return $ RealVal (sqrt x) pos
    builtin "log" [RealVal x pos] = return $ RealVal (log x) pos
    builtin "exp" [RealVal x pos] = return $ RealVal (exp x) pos
    builtin "op not" [LogVal b pos] = return $ LogVal (not b) pos
    builtin "op ~" [RealVal b pos] = return $ RealVal (-b) pos
    builtin fname _ = bad $ TypeError noLoc $ "Builtin " ++ fname

evalExp :: (Applicative m, Monad m) => Exp Type -> L0M m Value
evalExp (Literal val) = return val
evalExp (TupLit es _ pos) =
  TupVal <$> mapM evalExp es <*> pure pos
evalExp (ArrayLit es t pos) =
  arrayVal <$> mapM evalExp es <*> pure t <*> pure pos
evalExp (BinOp Plus e1 e2 (Int _) pos) = evalIntBinOp (+) e1 e2 pos
evalExp (BinOp Plus e1 e2 (Real _) pos) = evalRealBinOp (+) e1 e2 pos
evalExp (BinOp Minus e1 e2 (Int _) pos) = evalIntBinOp (-) e1 e2 pos
evalExp (BinOp Minus e1 e2 (Real _) pos) = evalRealBinOp (-) e1 e2 pos
evalExp (BinOp Pow e1 e2 (Int _) pos) = evalIntBinOp (^) e1 e2 pos
evalExp (BinOp Pow e1 e2 (Real _) pos) = evalRealBinOp (**) e1 e2 pos
evalExp (BinOp Times e1 e2 (Int _) pos) = evalIntBinOp (*) e1 e2 pos
evalExp (BinOp Times e1 e2 (Real _) pos) = evalRealBinOp (*) e1 e2 pos
evalExp (BinOp Divide e1 e2 (Int _) pos) = evalIntBinOp div e1 e2 pos
evalExp (BinOp Divide e1 e2 (Real _) pos) = evalRealBinOp (/) e1 e2 pos
evalExp (BinOp ShiftR e1 e2 _ pos) = evalIntBinOp shiftR e1 e2 pos
evalExp (BinOp ShiftL e1 e2 _ pos) = evalIntBinOp shiftL e1 e2 pos
evalExp (BinOp Band e1 e2 _ pos) = evalIntBinOp (.&.) e1 e2 pos
evalExp (BinOp Xor e1 e2 _ pos) = evalIntBinOp xor e1 e2 pos
evalExp (BinOp Bor e1 e2 _ pos) = evalIntBinOp (.|.) e1 e2 pos
evalExp (BinOp LogAnd e1 e2 _ pos) = evalBoolBinOp (&&) e1 e2 pos
evalExp (BinOp LogOr e1 e2 _ pos) = evalBoolBinOp (||) e1 e2 pos
evalExp (BinOp Equal e1 e2 _ pos) = do
  v1 <- evalExp e1
  v2 <- evalExp e2
  return $ LogVal (v1==v2) pos
evalExp (BinOp Less e1 e2 _ pos) = do
  v1 <- evalExp e1
  v2 <- evalExp e2
  return $ LogVal (v1<v2) pos
evalExp (BinOp Leq e1 e2 _ pos) = do
  v1 <- evalExp e1
  v2 <- evalExp e2
  return $ LogVal (v1<=v2) pos
evalExp (BinOp _ _ _ _ pos) = bad $ TypeError pos "evalExp Binop"
evalExp (And e1 e2 pos) = do
  v1 <- evalExp e1
  case v1 of LogVal True  _ -> evalExp e2
             LogVal False _ -> return v1
             _              -> bad $ TypeError pos "evalExp And"
evalExp (Or e1 e2 pos) = do
  v1 <- evalExp e1
  case v1 of LogVal True  _ -> return v1
             LogVal False _ -> evalExp e2
             _              -> bad $ TypeError pos "evalExp Or"
evalExp (Not e pos) = do
  v <- evalExp e
  case v of LogVal b _   -> return $ LogVal (not b) pos
            _            -> bad $ TypeError pos "evalExp Not"
evalExp (Negate e _ pos) = do
  v <- evalExp e
  case v of IntVal x _    -> return $ IntVal (-x) pos
            RealVal x _   -> return $ RealVal (-x) pos
            _             -> bad $ TypeError pos "evalExp Negate"
evalExp (If e1 e2 e3 _ pos) = do
  v <- evalExp e1
  case v of LogVal True _  -> evalExp e2
            LogVal False _ -> evalExp e3
            _              -> bad $ TypeError pos "evalExp If"
evalExp (Var name _ _) =
  lookupVar name
evalExp (Apply fname args _ _) = do
  fun <- lookupFun fname
  args' <- mapM evalExp args
  fun args'
evalExp (LetPat pat e body pos) = do
  v <- evalExp e
  case evalPattern pat v of
    Nothing   -> bad $ TypeError pos "evalExp Let pat"
    Just bnds -> local (`bindVars` bnds) $ evalExp body
evalExp (LetWith name e idxs ve body pos) = do
  v <- evalExp e
  idxs' <- mapM evalExp idxs
  vev <- evalExp ve
  v' <- change v idxs' vev
  binding [(name, v')] $ evalExp body
  where change _ [] to = return to
        change (ArrayVal arr t _) (IntVal i _:rest) to
          | i >= 0 && i <= upper = do
            let x = arr ! i
            x' <- change x rest to
            return $ ArrayVal (arr // [(i, x')]) t pos
          | otherwise = bad $ IndexOutOfBounds pos (upper+1) i
          where upper = snd $ bounds arr
        change _ _ _ = bad $ TypeError pos "evalExp Let Id"
evalExp (Index name idxs _ _ pos) = do
  v <- lookupVar name
  idxs' <- mapM evalExp idxs
  foldM idx v idxs'
  where idx (ArrayVal arr _ _) (IntVal i _)
          | i >= 0 && i <= upper = return $ arr ! i
          | otherwise             = bad $ IndexOutOfBounds pos (upper+1) i
          where upper = snd $ bounds arr
        idx _ _ = bad $ TypeError pos "evalExp Index"
evalExp (Iota e pos) = do
  v <- evalExp e
  case v of
    IntVal x _
      | x >= 0    -> return $ arrayVal (map (`IntVal` pos) [0..x-1]) (Int pos) pos
      | otherwise -> bad $ NegativeIota pos x
    _ -> bad $ TypeError pos "evalExp Iota"
evalExp (Size e pos) = do
  v <- evalExp e
  case v of
    ArrayVal arr _ _ -> return $ IntVal (snd $ bounds arr) pos
    _ -> bad $ TypeError pos "evalExp Size"
evalExp (Replicate e1 e2 et pos) = do
  v1 <- evalExp e1
  v2 <- evalExp e2
  case v1 of
    IntVal x _
      | x >= 0    -> return $ ArrayVal (listArray (0,x-1) $ repeat v2) et pos
      | otherwise -> bad $ NegativeReplicate pos x
    _   -> bad $ TypeError pos "evalExp Replicate"
evalExp (Reshape shapeexp arrexp _ outtype pos) = do
  shape <- mapM (asInt <=< evalExp) shapeexp
  arr <- evalExp arrexp
  let reshape (Array t _ _) (n:rest) vs
        | length vs `mod` n == 0 =
          arrayVal <$> (mapM (reshape t rest) (chunk (length vs `div` n) vs))
                   <*> pure t <*> pure pos
        | otherwise = bad $ InvalidArrayShape pos (arrayShape arr) shape
      reshape _ [] [v] = return v
      reshape _ _ _ = bad $ TypeError pos "evalExp Reshape reshape"
  reshape outtype shape $ flatten arr
  where flatten (ArrayVal arr _ _) = concatMap flatten $ elems arr
        flatten t = [t]
        chunk _ [] = []
        chunk i l = let (a,b) = splitAt i l
                    in a : chunk i b
        asInt (IntVal x _) = return x
        asInt _ = bad $ TypeError pos "evalExp Reshape int"
evalExp (Transpose arrexp _ _ pos) = do
  v <- evalExp arrexp
  case v of
    ArrayVal inarr (Array et _ _) _ -> do
      let arr el = arrayVal el et pos
      els' <- map arr <$> transpose <$> mapM arrToList (elems inarr)
      return $ arrayVal els' (Array et Nothing pos) pos
    _ -> bad $ TypeError pos "evalExp Transpose"
evalExp (Map fun e _ outtype pos) = do
  vs <- arrToList =<< evalExp e
  case outtype of
    (Array t _ _) -> do<commit_msg>Remove superfluous parentheses in interpreter.<commit_after>{-# LANGUAGE GeneralizedNewtypeDeriving, ScopedTypeVariables #-}
module L0.Interpreter
  ( runProg
  , runProgIO
  , InterpreterError(..) )
where

import Control.Applicative
import Control.Monad.Reader
import Control.Monad.Trans.Either

import Data.Array
import Data.Bits
import Data.List
import Data.Loc
import qualified Data.Map as M

import System.IO (hFlush, stdout)

import L0.AbSyn
import L0.Parser

data InterpreterError = MissingMainFunction
                      | IndexOutOfBounds Loc Int Int
                      -- ^ First @Int@ is array size, second is attempted index.
                      | NegativeIota Loc Int
                      | NegativeReplicate Loc Int
                      | ReadError Loc Type String
                      | InvalidArrayShape Loc [Int] [Int]
                      -- ^ First @Int@ is old shape, second is attempted new shape.
                      | TypeError Loc String

instance Show InterpreterError where
  show MissingMainFunction =
    "No main function defined."
  show (IndexOutOfBounds pos arrsz i) =
    "Array index " ++ show i ++ " out of bounds of array size " ++ show arrsz ++ " at " ++ locStr pos ++ "."
  show (NegativeIota pos n) =
    "Argument " ++ show n ++ " to iota at " ++ locStr pos ++ " is negative."
  show (NegativeReplicate pos n) =
    "Argument " ++ show n ++ " to replicate at " ++ locStr pos ++ " is negative."
  show (TypeError pos s) =
    "Type error at " ++ locStr pos ++ " in " ++ s ++ " during interpretation.  This implies a bug in the type checker."
  show (ReadError pos t s) =
    "Read error while trying to read " ++ ppType t ++ " at " ++ locStr pos ++ ".  Input line was: " ++ s
  show (InvalidArrayShape pos shape newshape) =
    "Invalid array reshaping at " ++ locStr pos ++ ", from " ++ show shape ++ " to " ++ show newshape


data L0Env m = L0Env { envVtable  :: M.Map String Value
                     , envFtable  :: M.Map String ([Value] -> L0M m Value)
                     , envWriteOp :: String -> L0M m ()
                     , envReadOp  :: L0M m String }

newtype L0M m a = L0M (ReaderT (L0Env m) (EitherT InterpreterError m) a)
  deriving (MonadReader (L0Env m), Monad, Applicative, Functor)

runL0M :: L0M m a -> L0Env m -> m (Either InterpreterError a)
runL0M (L0M m) env = runEitherT (runReaderT m env)

bad :: Monad m => InterpreterError -> L0M m a
bad = L0M . lift . left

bindVar :: L0Env m -> (String, Value) -> L0Env m
bindVar env (name,val) =
  env { envVtable = M.insert name val $ envVtable env }

bindVars :: L0Env m -> [(String, Value)] -> L0Env m
bindVars = foldl bindVar

binding :: Monad m => [(String, Value)] -> L0M m a -> L0M m a
binding bnds = local (`bindVars` bnds)

lookupVar :: Monad m => String -> L0M m Value
lookupVar fname = do val <- asks $ M.lookup fname . envVtable
                     case val of Just val' -> return val'
                                 Nothing   -> bad $ TypeError noLoc "lookupVar"

lookupFun :: Monad m => String -> L0M m ([Value] -> L0M m Value)
lookupFun fname = do fun <- asks $ M.lookup fname . envFtable
                     case fun of Just fun' -> return fun'
                                 Nothing   -> bad $ TypeError noLoc "lookupFun"

arrToList :: Monad m => Value -> L0M m [Value]
arrToList (ArrayVal l _ _) = return $ elems l
arrToList v = bad $ TypeError (locOf v) "arrToList"

tupToList :: Monad m => Value -> L0M m [Value]
tupToList (TupVal l _) = return l
tupToList v = bad $ TypeError (locOf v) "tupToList"

runProgIO :: Prog Type -> IO (Either InterpreterError Value)
runProgIO = runProg putStr (hFlush stdout >> getLine)

runProg :: (Applicative m, Monad m) => (String -> m ()) -> m String
        -> Prog Type -> m (Either InterpreterError Value)
runProg wop rop prog = do
  let ftable = foldl expand builtins prog
      l0env = L0Env { envVtable = M.empty
                    , envFtable = ftable
                    , envWriteOp = L0M . lift . lift . wop
                    , envReadOp = L0M $ lift $ lift rop }
      runmain = case M.lookup "main" ftable of
                  Nothing -> bad MissingMainFunction
                  Just mainfun -> mainfun []
  runL0M runmain l0env
  where
    -- We assume that the program already passed the type checker, so
    -- we don't check for duplicate definitions or anything.
    expand ftable (name,_,params,body,_) =
        let fun args = binding (zip (map fst params) args) $ evalExp body
        in M.insert name fun ftable
    builtins = M.fromList [("toReal", builtin "toReal")
                          ,("trunc", builtin "trunc")
                          ,("sqrt", builtin "sqrt")
                          ,("log", builtin "log")
                          ,("exp", builtin "exp")
                          ,("op not", builtin "op not")
                          ,("op ~", builtin "op ~")]
    builtin "toReal" [IntVal x pos] = return $ RealVal (fromIntegral x) pos
    builtin "trunc" [RealVal x pos] = return $ IntVal (truncate x) pos
    builtin "sqrt" [RealVal x pos] = return $ RealVal (sqrt x) pos
    builtin "log" [RealVal x pos] = return $ RealVal (log x) pos
    builtin "exp" [RealVal x pos] = return $ RealVal (exp x) pos
    builtin "op not" [LogVal b pos] = return $ LogVal (not b) pos
    builtin "op ~" [RealVal b pos] = return $ RealVal (-b) pos
    builtin fname _ = bad $ TypeError noLoc $ "Builtin " ++ fname

evalExp :: (Applicative m, Monad m) => Exp Type -> L0M m Value
evalExp (Literal val) = return val
evalExp (TupLit es _ pos) =
  TupVal <$> mapM evalExp es <*> pure pos
evalExp (ArrayLit es t pos) =
  arrayVal <$> mapM evalExp es <*> pure t <*> pure pos
evalExp (BinOp Plus e1 e2 (Int _) pos) = evalIntBinOp (+) e1 e2 pos
evalExp (BinOp Plus e1 e2 (Real _) pos) = evalRealBinOp (+) e1 e2 pos
evalExp (BinOp Minus e1 e2 (Int _) pos) = evalIntBinOp (-) e1 e2 pos
evalExp (BinOp Minus e1 e2 (Real _) pos) = evalRealBinOp (-) e1 e2 pos
evalExp (BinOp Pow e1 e2 (Int _) pos) = evalIntBinOp (^) e1 e2 pos
evalExp (BinOp Pow e1 e2 (Real _) pos) = evalRealBinOp (**) e1 e2 pos
evalExp (BinOp Times e1 e2 (Int _) pos) = evalIntBinOp (*) e1 e2 pos
evalExp (BinOp Times e1 e2 (Real _) pos) = evalRealBinOp (*) e1 e2 pos
evalExp (BinOp Divide e1 e2 (Int _) pos) = evalIntBinOp div e1 e2 pos
evalExp (BinOp Divide e1 e2 (Real _) pos) = evalRealBinOp (/) e1 e2 pos
evalExp (BinOp ShiftR e1 e2 _ pos) = evalIntBinOp shiftR e1 e2 pos
evalExp (BinOp ShiftL e1 e2 _ pos) = evalIntBinOp shiftL e1 e2 pos
evalExp (BinOp Band e1 e2 _ pos) = evalIntBinOp (.&.) e1 e2 pos
evalExp (BinOp Xor e1 e2 _ pos) = evalIntBinOp xor e1 e2 pos
evalExp (BinOp Bor e1 e2 _ pos) = evalIntBinOp (.|.) e1 e2 pos
evalExp (BinOp LogAnd e1 e2 _ pos) = evalBoolBinOp (&&) e1 e2 pos
evalExp (BinOp LogOr e1 e2 _ pos) = evalBoolBinOp (||) e1 e2 pos
evalExp (BinOp Equal e1 e2 _ pos) = do
  v1 <- evalExp e1
  v2 <- evalExp e2
  return $ LogVal (v1==v2) pos
evalExp (BinOp Less e1 e2 _ pos) = do
  v1 <- evalExp e1
  v2 <- evalExp e2
  return $ LogVal (v1<v2) pos
evalExp (BinOp Leq e1 e2 _ pos) = do
  v1 <- evalExp e1
  v2 <- evalExp e2
  return $ LogVal (v1<=v2) pos
evalExp (BinOp _ _ _ _ pos) = bad $ TypeError pos "evalExp Binop"
evalExp (And e1 e2 pos) = do
  v1 <- evalExp e1
  case v1 of LogVal True  _ -> evalExp e2
             LogVal False _ -> return v1
             _              -> bad $ TypeError pos "evalExp And"
evalExp (Or e1 e2 pos) = do
  v1 <- evalExp e1
  case v1 of LogVal True  _ -> return v1
             LogVal False _ -> evalExp e2
             _              -> bad $ TypeError pos "evalExp Or"
evalExp (Not e pos) = do
  v <- evalExp e
  case v of LogVal b _   -> return $ LogVal (not b) pos
            _            -> bad $ TypeError pos "evalExp Not"
evalExp (Negate e _ pos) = do
  v <- evalExp e
  case v of IntVal x _    -> return $ IntVal (-x) pos
            RealVal x _   -> return $ RealVal (-x) pos
            _             -> bad $ TypeError pos "evalExp Negate"
evalExp (If e1 e2 e3 _ pos) = do
  v <- evalExp e1
  case v of LogVal True _  -> evalExp e2
            LogVal False _ -> evalExp e3
            _              -> bad $ TypeError pos "evalExp If"
evalExp (Var name _ _) =
  lookupVar name
evalExp (Apply fname args _ _) = do
  fun <- lookupFun fname
  args' <- mapM evalExp args
  fun args'
evalExp (LetPat pat e body pos) = do
  v <- evalExp e
  case evalPattern pat v of
    Nothing   -> bad $ TypeError pos "evalExp Let pat"
    Just bnds -> local (`bindVars` bnds) $ evalExp body
evalExp (LetWith name e idxs ve body pos) = do
  v <- evalExp e
  idxs' <- mapM evalExp idxs
  vev <- evalExp ve
  v' <- change v idxs' vev
  binding [(name, v')] $ evalExp body
  where change _ [] to = return to
        change (ArrayVal arr t _) (IntVal i _:rest) to
          | i >= 0 && i <= upper = do
            let x = arr ! i
            x' <- change x rest to
            return $ ArrayVal (arr // [(i, x')]) t pos
          | otherwise = bad $ IndexOutOfBounds pos (upper+1) i
          where upper = snd $ bounds arr
        change _ _ _ = bad $ TypeError pos "evalExp Let Id"
evalExp (Index name idxs _ _ pos) = do
  v <- lookupVar name
  idxs' <- mapM evalExp idxs
  foldM idx v idxs'
  where idx (ArrayVal arr _ _) (IntVal i _)
          | i >= 0 && i <= upper = return $ arr ! i
          | otherwise             = bad $ IndexOutOfBounds pos (upper+1) i
          where upper = snd $ bounds arr
        idx _ _ = bad $ TypeError pos "evalExp Index"
evalExp (Iota e pos) = do
  v <- evalExp e
  case v of
    IntVal x _
      | x >= 0    -> return $ arrayVal (map (`IntVal` pos) [0..x-1]) (Int pos) pos
      | otherwise -> bad $ NegativeIota pos x
    _ -> bad $ TypeError pos "evalExp Iota"
evalExp (Size e pos) = do
  v <- evalExp e
  case v of
    ArrayVal arr _ _ -> return $ IntVal (snd $ bounds arr) pos
    _ -> bad $ TypeError pos "evalExp Size"
evalExp (Replicate e1 e2 et pos) = do
  v1 <- evalExp e1
  v2 <- evalExp e2
  case v1 of
    IntVal x _
      | x >= 0    -> return $ ArrayVal (listArray (0,x-1) $ repeat v2) et pos
      | otherwise -> bad $ NegativeReplicate pos x
    _   -> bad $ TypeError pos "evalExp Replicate"
evalExp (Reshape shapeexp arrexp _ outtype pos) = do
  shape <- mapM (asInt <=< evalExp) shapeexp
  arr <- evalExp arrexp
  let reshape (Array t _ _) (n:rest) vs
        | length vs `mod` n == 0 =
          arrayVal <$> mapM (reshape t rest) (chunk (length vs `div` n) vs)
                   <*> pure t <*> pure pos
        | otherwise = bad $ InvalidArrayShape pos (arrayShape arr) shape
      reshape _ [] [v] = return v
      reshape _ _ _ = bad $ TypeError pos "evalExp Reshape reshape"
  reshape outtype shape $ flatten arr
  where flatten (ArrayVal arr _ _) = concatMap flatten $ elems arr
        flatten t = [t]
        chunk _ [] = []
        chunk i l = let (a,b) = splitAt i l
                    in a : chunk i b
        asInt (IntVal x _) = return x
        asInt _ = bad $ TypeError pos "evalExp Reshape int"
evalExp (Transpose arrexp _ _ pos) = do
  v <- evalExp arrexp
  case v of
    ArrayVal inarr (Array et _ _) _ -> do
      let arr el = arrayVal el et pos
      els' <- map arr <$> transpose <$> mapM arrToList (elems inarr)
      return $ arrayVal els' (Array et Nothing pos) pos
    _ -> bad $ TypeError pos "evalExp Transpose"
evalExp (Map fun e _ outtype pos) = do
  vs <- arrToList =<< evalExp e
  case outtype of
    (Array t _ _) -> do
{'fraction_non_alphanumeric': 0.10377730117739747, 'fraction_numerical': 0.014706526020921784, 'mean_word_length': 2.227118644067797, 'pattern_counts': {'":': 0, '<': 103, '<?xml version=': 0, '>': 161, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '77040', 'n_tokens_mistral': 7811, 'n_tokens_neox': 7621, 'n_words': 2921}