path: root/src/Interpreter.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImplicitParams #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
module Interpreter (
  interpret,
  interpretOpen,
  Value(..),
) where

import Control.Monad (foldM, join, when, forM_)
import Data.Bitraversable (bitraverse)
import Data.Char (isSpace)
import Data.Functor.Identity
import Data.Int (Int64)
import Data.IORef
import System.IO (hPutStrLn, stderr)
import System.IO.Unsafe (unsafePerformIO)

import Debug.Trace

import Array
import AST
import AST.Pretty
import CHAD.Types
import Data
import Interpreter.Rep


newtype AcM s a = AcM { unAcM :: IO a }
  deriving newtype (Functor, Applicative, Monad)

runAcM :: (forall s. AcM s a) -> a
runAcM (AcM m) = unsafePerformIO m

acmDebugLog :: String -> AcM s ()
acmDebugLog s = AcM (hPutStrLn stderr s)

interpret :: Ex '[] t -> Rep t
interpret = interpretOpen False SNil

-- | Bool: whether to trace execution with debug prints (very verbose)
interpretOpen :: Bool -> SList Value env -> Ex env t -> Rep t
interpretOpen prints env e =
  runAcM $
    let ?depth = 0
        ?prints = prints
    in interpret' env e

interpret' :: forall env t s. (?prints :: Bool, ?depth :: Int) => SList Value env -> Ex env t -> AcM s (Rep t)
interpret' env e = do
  let dep = ?depth
  let lenlimit = max 20 (100 - dep)
  let replace a b = map (\c -> if c == a then b else c)
  let trunc s | length s > lenlimit = take (lenlimit - 3) (replace '\n' ' ' s) ++ "..."
              | otherwise           = replace '\n' ' ' s
  when ?prints $ acmDebugLog $ replicate dep ' ' ++ "ev: " ++ trunc (ppExpr env e)
  res <- let ?depth = dep + 1 in interpret'Rec env e
  when ?prints $ acmDebugLog $ replicate dep ' ' ++ "<- " ++ showValue 0 (typeOf e) res ""
  return res

interpret'Rec :: forall env t s. (?prints :: Bool, ?depth :: Int) => SList Value env -> Ex env t -> AcM s (Rep t)
interpret'Rec env = \case
  EVar _ _ i -> case slistIdx env i of Value x -> return x
  ELet _ a b -> do
    x <- interpret' env a
    let ?depth = ?depth - 1 in interpret' (Value x `SCons` env) b
  expr | False && trace ("<i> " ++ takeWhile (not . isSpace) (show expr)) False -> undefined
  EPair _ a b -> (,) <$> interpret' env a <*> interpret' env b
  EFst _ e -> fst <$> interpret' env e
  ESnd _ e -> snd <$> interpret' env e
  ENil _ -> return ()
  EInl _ _ e -> Left <$> interpret' env e
  EInr _ _ e -> Right <$> interpret' env e
  ECase _ e a b -> interpret' env e >>= \case
                     Left x -> interpret' (Value x `SCons` env) a
                     Right y -> interpret' (Value y `SCons` env) b
  ENothing _ _ -> return Nothing
  EJust _ e -> Just <$> interpret' env e
  EMaybe _ a b e -> maybe (interpret' env a) (\x -> interpret' (Value x `SCons` env) b) =<< interpret' env e
  EConstArr _ _ _ v -> return v
  EBuild _ dim a b -> do
    sh <- unTupRepIdx ShNil ShCons dim <$> interpret' env a
    arrayGenerateM sh (\idx -> interpret' (Value (tupRepIdx ixUncons dim idx) `SCons` env) b)
  EFold1Inner _ a b c -> do
    let f = \x y -> interpret' (Value y `SCons` Value x `SCons` env) a
    x0 <- interpret' env b
    arr <- interpret' env c
    let sh `ShCons` n = arrayShape arr
    arrayGenerateM sh $ \idx -> foldM f x0 [arrayIndex arr (idx `IxCons` i) | i <- [0 .. n - 1]]
  ESum1Inner _ e -> do
    arr <- interpret' env e
    let STArr _ (STScal t) = typeOf e
        sh `ShCons` n = arrayShape arr
    numericIsNum t $ return $ arrayGenerate sh $ \idx -> sum [arrayIndex arr (idx `IxCons` i) | i <- [0 .. n - 1]]
  EUnit _ e -> arrayGenerateLinM ShNil (\_ -> interpret' env e)
  EReplicate1Inner _ a b -> do
    n <- fromIntegral @Int64 @Int <$> interpret' env a
    arr <- interpret' env b
    let sh = arrayShape arr
    return $ arrayGenerate (sh `ShCons` n) (\(idx `IxCons` _) -> arrayIndex arr idx)
  EMaximum1Inner _ e -> do
    arr <- interpret' env e
    let STArr _ (STScal t) = typeOf e
        sh `ShCons` n = arrayShape arr
    numericIsNum t $ return $
      arrayGenerate sh (\idx -> maximum [arrayIndex arr (idx `IxCons` i) | i <- [0 .. n-1]])
  EMinimum1Inner _ e -> do
    arr <- interpret' env e
    let STArr _ (STScal t) = typeOf e
        sh `ShCons` n = arrayShape arr
    numericIsNum t $ return $
      arrayGenerate sh (\idx -> minimum [arrayIndex arr (idx `IxCons` i) | i <- [0 .. n-1]])
  EConst _ _ v -> return v
  EIdx0 _ e -> (`arrayIndexLinear` 0) <$> interpret' env e
  EIdx1 _ a b -> arrayIndex1 <$> interpret' env a <*> (fromIntegral @Int64 @Int <$> interpret' env b)
  EIdx _ a b
    | STArr n _ <- typeOf a
    -> arrayIndex <$> interpret' env a <*> (unTupRepIdx IxNil IxCons n <$> interpret' env b)
  EShape _ e | STArr n _ <- typeOf e -> tupRepIdx shUncons n . arrayShape <$> interpret' env e
  EOp _ op e -> interpretOp op <$> interpret' env e
  ECustom _ _ _ _ pr _ _ e1 e2 -> do
    e1' <- interpret' env e1
    e2' <- interpret' env e2
    interpret' (Value e2' `SCons` Value e1' `SCons` SNil) pr
  EWith _ t e1 e2 -> do
    initval <- interpret' env e1
    withAccum t (typeOf e2) initval $ \accum ->
      interpret' (Value accum `SCons` env) e2
  EAccum _ t p e1 e2 e3 -> do
    idx <- interpret' env e1
    val <- interpret' env e2
    accum <- interpret' env e3
    accumAddSparse t p accum idx val
  EZero _ t -> do
    return $ zeroD2 t
  EPlus _ t a b -> do
    a' <- interpret' env a
    b' <- interpret' env b
    return $ addD2s t a' b'
  EOneHot _ t p a b -> do
    a' <- interpret' env a
    b' <- interpret' env b
    return $ onehotD2 p t a' b'
  EError _ _ s -> error $ "Interpreter: Program threw error: " ++ s

interpretOp :: SOp a t -> Rep a -> Rep t
interpretOp op arg = case op of
  OAdd st -> numericIsNum st $ uncurry (+) arg
  OMul st -> numericIsNum st $ uncurry (*) arg
  ONeg st -> numericIsNum st $ negate arg
  OLt st -> numericIsNum st $ uncurry (<) arg
  OLe st -> numericIsNum st $ uncurry (<=) arg
  OEq st -> styIsEq st $ uncurry (==) arg
  ONot -> not arg
  OAnd -> uncurry (&&) arg
  OOr -> uncurry (||) arg
  OIf -> if arg then Left () else Right ()
  ORound64 -> round arg
  OToFl64 -> fromIntegral arg
  ORecip st -> floatingIsFractional st $ recip arg
  OExp st -> floatingIsFractional st $ exp arg
  OLog st -> floatingIsFractional st $ log arg
  OIDiv st -> integralIsIntegral st $ uncurry quot arg
  where
    styIsEq :: SScalTy t -> (Eq (Rep (TScal t)) => r) -> r
    styIsEq STI32 = id
    styIsEq STI64 = id
    styIsEq STF32 = id
    styIsEq STF64 = id
    styIsEq STBool = id

zeroD2 :: STy t -> Rep (D2 t)
zeroD2 typ = case typ of
  STNil -> ()
  STPair _ _ -> Nothing
  STEither _ _ -> Nothing
  STMaybe _ -> Nothing
  STArr SZ t -> arrayUnit (zeroD2 t)
  STArr n _ -> emptyArray n
  STScal sty -> case sty of
                  STI32 -> ()
                  STI64 -> ()
                  STF32 -> 0.0
                  STF64 -> 0.0
                  STBool -> ()
  STAccum{} -> error "Zero of Accum"

addD2s :: STy t -> Rep (D2 t) -> Rep (D2 t) -> Rep (D2 t)
addD2s typ a b = case typ of
  STNil -> ()
  STPair t1 t2 -> case (a, b) of
    (Nothing, _) -> b
    (_, Nothing) -> a
    (Just (x1, x2), Just (y1, y2)) -> Just (addD2s t1 x1 y1, addD2s t2 x2 y2)
  STEither t1 t2 -> case (a, b) of
    (Nothing, _) -> b
    (_, Nothing) -> a
    (Just (Left x), Just (Left y)) -> Just (Left (addD2s t1 x y))
    (Just (Right x), Just (Right y)) -> Just (Right (addD2s t2 x y))
    _ -> error "Plus of inconsistent Eithers"
  STMaybe t -> case (a, b) of
    (Nothing, _) -> b
    (_, Nothing) -> a
    (Just x, Just y) -> Just (addD2s t x y)
  STArr _ t ->
    let sh1 = arrayShape a
        sh2 = arrayShape b
    in if | shapeSize sh1 == 0 -> b
          | shapeSize sh2 == 0 -> a
          | sh1 == sh2 -> arrayGenerateLin sh1 (\i -> addD2s t (arrayIndexLinear a i) (arrayIndexLinear b i))
          | otherwise -> error "Plus of inconsistently shaped arrays"
  STScal sty -> case sty of
    STI32 -> ()
    STI64 -> ()
    STF32 -> a + b
    STF64 -> a + b
    STBool -> ()
  STAccum{} -> error "Plus of Accum"

onehotD2 :: SAcPrj p a b -> STy a -> Rep (AcIdx p a) -> Rep (D2 b) -> Rep (D2 a)
onehotD2 SAPHere _ _ val = val
onehotD2 (SAPFst prj) (STPair a b) idx val = Just (onehotD2 prj a idx val, zeroD2 b)
onehotD2 (SAPSnd prj) (STPair a b) idx val = Just (zeroD2 a, onehotD2 prj b idx val)
onehotD2 (SAPLeft prj) (STEither a _) idx val = Just (Left (onehotD2 prj a idx val))
onehotD2 (SAPRight prj) (STEither _ b) idx val = Just (Right (onehotD2 prj b idx val))
onehotD2 (SAPJust prj) (STMaybe a) idx val = Just (onehotD2 prj a idx val)
onehotD2 (SAPArrIdx prj _) (STArr n a) idx val =
  runIdentity $ onehotArray (\idx' -> Identity (onehotD2 prj a idx' val)) (Identity (zeroD2 a)) n prj idx

withAccum :: STy t -> STy a -> Rep (D2 t) -> (RepAc t -> AcM s (Rep a)) -> AcM s (Rep a, Rep (D2 t))
withAccum t _ initval f = AcM $ do
  accum <- newAcSparse t SAPHere () initval
  out <- unAcM $ f accum
  val <- readAcSparse t accum
  return (out, val)

newAcZero :: STy t -> IO (RepAc t)
newAcZero = \case
  STNil -> return ()
  STPair{} -> newIORef Nothing
  STEither{} -> newIORef Nothing
  STMaybe _ -> newIORef Nothing
  STArr n _ -> newIORef (emptyArray n)
  STScal sty -> case sty of
    STI32 -> return ()
    STI64 -> return ()
    STF32 -> newIORef 0.0
    STF64 -> newIORef 0.0
    STBool -> return ()
  STAccum{} -> error "Nested accumulators"

newAcSparse :: STy a -> SAcPrj p a b -> Rep (AcIdx p a) -> Rep (D2 b) -> IO (RepAc a)
newAcSparse typ prj idx val = case (typ, prj) of
  (STNil, SAPHere) -> return ()
  (STPair t1 t2, SAPHere) -> newIORef =<< traverse (bitraverse (newAcSparse t1 SAPHere ()) (newAcSparse t2 SAPHere ())) val
  (STEither t1 t2, SAPHere) -> newIORef =<< traverse (bitraverse (newAcSparse t1 SAPHere ()) (newAcSparse t2 SAPHere ())) val
  (STMaybe t1, SAPHere) -> newIORef =<< traverse (newAcSparse t1 SAPHere ()) val
  (STArr _ t1, SAPHere) -> newIORef =<< traverse (newAcSparse t1 SAPHere ()) val
  (STScal sty, SAPHere) -> case sty of
    STI32 -> return ()
    STI64 -> return ()
    STF32 -> newIORef val
    STF64 -> newIORef val
    STBool -> return ()

  (STPair t1 t2, SAPFst prj') ->
    newIORef . Just =<< (,) <$> newAcSparse t1 prj' idx val <*> newAcZero t2
  (STPair t1 t2, SAPSnd prj') ->
    newIORef . Just =<< (,) <$> newAcZero t1 <*> newAcSparse t2 prj' idx val

  (STEither t1 _, SAPLeft prj') -> newIORef . Just . Left =<< newAcSparse t1 prj' idx val
  (STEither _ t2, SAPRight prj') -> newIORef . Just . Right =<< newAcSparse t2 prj' idx val

  (STMaybe t1, SAPJust prj') -> newIORef . Just =<< newAcSparse t1 prj' idx val

  (STArr n t, SAPArrIdx prj' _) -> newIORef =<< newAcArray n t prj' idx val

  (STAccum{}, _) -> error "Accumulators not allowed in source program"

newAcArray :: SNat n -> STy a -> SAcPrj p a b -> Rep (AcIdx (APArrIdx p) (TArr n a)) -> Rep (D2 b) -> IO (Array n (RepAc a))
newAcArray n t prj idx val = onehotArray (\idx' -> newAcSparse t prj idx' val) (newAcZero t) n prj idx

onehotArray :: Monad m
            => (Rep (AcIdx p a) -> m v)  -- ^ the "one"
            -> m v  -- ^ the "zero"
            -> SNat n -> SAcPrj p a b -> Rep (AcIdx (APArrIdx p) (TArr n a)) -> m (Array n v)
onehotArray mkone mkzero n _ ((arrindex', arrsh'), idx) =
  let arrindex = unTupRepIdx IxNil IxCons n arrindex'
      arrsh = unTupRepIdx ShNil ShCons n arrsh'
      !linindex = toLinearIndex arrsh arrindex
  in arrayGenerateLinM arrsh (\i -> if i == linindex then mkone idx else mkzero)

readAcSparse :: STy t -> RepAc t -> IO (Rep (D2 t))
readAcSparse typ val = case typ of
  STNil -> return ()
  STPair t1 t2 -> traverse (bitraverse (readAcSparse t1) (readAcSparse t2)) =<< readIORef val
  STEither t1 t2 -> traverse (bitraverse (readAcSparse t1) (readAcSparse t2)) =<< readIORef val
  STMaybe t -> traverse (readAcSparse t) =<< readIORef val
  STArr _ t -> traverse (readAcSparse t) =<< readIORef val
  STScal sty -> case sty of
    STI32 -> return ()
    STI64 -> return ()
    STF32 -> readIORef val
    STF64 -> readIORef val
    STBool -> return ()
  STAccum{} -> error "Nested accumulators"

accumAddSparse :: STy a -> SAcPrj p a b -> RepAc a -> Rep (AcIdx p a) -> Rep (D2 b) -> AcM s ()
accumAddSparse typ prj ref idx val = case (typ, prj) of
  (STNil, SAPHere) -> return ()

  (STPair t1 t2, SAPHere) ->
    case val of
      Nothing -> return ()
      Just (val1, val2) ->
        realiseMaybeSparse ref ((,) <$> newAcSparse t1 SAPHere () val1
                                    <*> newAcSparse t2 SAPHere () val2)
                               (\(ac1, ac2) -> do accumAddSparse t1 SAPHere ac1 () val1
                                                  accumAddSparse t2 SAPHere ac2 () val2)
  (STPair t1 t2, SAPFst prj') ->
    realiseMaybeSparse ref ((,) <$> newAcSparse t1 prj' idx val <*> newAcZero t2)
                           (\(ac1, _) -> do accumAddSparse t1 prj' ac1 idx val)
  (STPair t1 t2, SAPSnd prj') ->
    realiseMaybeSparse ref ((,) <$> newAcZero t1 <*> newAcSparse t2 prj' idx val)
                           (\(_, ac2) -> do accumAddSparse t2 prj' ac2 idx val)

  (STEither{}, SAPHere) ->
    case val of
      Nothing -> return ()
      Just (Left val1) -> accumAddSparse typ (SAPLeft SAPHere) ref () val1
      Just (Right val2) -> accumAddSparse typ (SAPRight SAPHere) ref () val2
  (STEither t1 _, SAPLeft prj') ->
    realiseMaybeSparse ref (Left <$> newAcSparse t1 prj' idx val)
                           (\case Left ac1 -> accumAddSparse t1 prj' ac1 idx val
                                  Right{} -> error "Mismatched Either in accumAddSparse (r +l)")
  (STEither _ t2, SAPRight prj') ->
    realiseMaybeSparse ref (Right <$> newAcSparse t2 prj' idx val)
                           (\case Right ac2 -> accumAddSparse t2 prj' ac2 idx val
                                  Left{} -> error "Mismatched Either in accumAddSparse (l +r)")

  (STMaybe{}, SAPHere) ->
    case val of
      Nothing -> return ()
      Just val' -> accumAddSparse typ (SAPJust SAPHere) ref () val'
  (STMaybe t1, SAPJust prj') ->
      realiseMaybeSparse ref (newAcSparse t1 prj' idx val)
                             (\ac -> accumAddSparse t1 prj' ac idx val)

  (STArr _ t1, SAPHere) ->
    let add ac = forM_ [0 .. arraySize ac - 1] $ \i ->
                   unAcM $ accumAddSparse t1 SAPHere (arrayIndexLinear ac i) () (arrayIndexLinear val i)
    in if arraySize val == 0
         then return ()
         else AcM $ join $ atomicModifyIORef' ref $ \ac ->
                if arraySize ac == 0
                  then (ac, do newac <- arrayMapM (newAcSparse t1 SAPHere ()) val
                               join $ atomicModifyIORef' ref $ \ac' ->
                                 if arraySize ac == 0
                                   then (newac, return ())
                                   else (ac', add ac'))
                  else (ac, add ac)
  (STArr n t1, SAPArrIdx prj' _) ->
    let ((arrindex', arrsh'), idx') = idx
        arrindex = unTupRepIdx IxNil IxCons n arrindex'
        arrsh = unTupRepIdx ShNil ShCons n arrsh'
        linindex = toLinearIndex arrsh arrindex
        add ac = unAcM $ accumAddSparse t1 prj' (arrayIndexLinear ac linindex) idx' val
    in AcM $ join $ atomicModifyIORef' ref $ \ac ->
         if arraySize ac == 0
           then (ac, do newac <- onehotArray (\_ -> newAcSparse t1 prj' idx' val) (newAcZero t1) n prj' idx
                        join $ atomicModifyIORef' ref $ \ac' ->
                          if arraySize ac == 0
                            then (newac, return ())
                            else (ac', add ac'))
           else (ac, add ac)

  (STScal sty, SAPHere) -> AcM $ case sty of
    STI32 -> return ()
    STI64 -> return ()
    STF32 -> atomicModifyIORef' ref (\x -> (x + val, ()))
    STF64 -> atomicModifyIORef' ref (\x -> (x + val, ()))
    STBool -> return ()

  (STAccum{}, _) -> error "Accumulators not allowed in source program"

realiseMaybeSparse :: IORef (Maybe a) -> IO a -> (a -> AcM s ()) -> AcM s ()
realiseMaybeSparse ref makeval modifyval =
  -- Try modifying what's already in ref. The 'join' makes the snd
  -- of the function's return value a _continuation_ that is run after
  -- the critical section ends.
  AcM $ join $ atomicModifyIORef' ref $ \ac -> case ac of
    -- Oops, ref's contents was still sparse. Have to initialise
    -- it first, then try again.
    Nothing -> (ac, do val <- makeval
                       join $ atomicModifyIORef' ref $ \ac' -> case ac' of
                                Nothing -> (Just val, return ())
                                Just val' -> (ac', unAcM $ modifyval val'))
    -- Yep, ref already had a value in there, so we can just add
    -- val' to it recursively.
    Just val -> (ac, unAcM $ modifyval val)


numericIsNum :: ScalIsNumeric st ~ True => SScalTy st -> ((Num (ScalRep st), Ord (ScalRep st)) => r) -> r
numericIsNum STI32 = id
numericIsNum STI64 = id
numericIsNum STF32 = id
numericIsNum STF64 = id

floatingIsFractional :: ScalIsFloating st ~ True => SScalTy st -> ((Floating (ScalRep st), Ord (ScalRep st), ScalIsNumeric st ~ True, ScalIsFloating st ~ True) => r) -> r
floatingIsFractional STF32 = id
floatingIsFractional STF64 = id

integralIsIntegral :: ScalIsIntegral st ~ True => SScalTy st -> ((Integral (ScalRep st), Ord (ScalRep st), ScalIsNumeric st ~ True, ScalIsIntegral st ~ True) => r) -> r
integralIsIntegral STI32 = id
integralIsIntegral STI64 = id

unTupRepIdx :: f Z -> (forall m. f m -> Int -> f (S m))
            -> SNat n -> Rep (Tup (Replicate n TIx)) -> f n
unTupRepIdx nil _    SZ     _        = nil
unTupRepIdx nil cons (SS n) (idx, i) = unTupRepIdx nil cons n idx `cons` fromIntegral @Int64 @Int i

tupRepIdx :: (forall m. f (S m) -> (f m, Int))
          -> SNat n -> f n -> Rep (Tup (Replicate n TIx))
tupRepIdx _      SZ     _   = ()
tupRepIdx uncons (SS n) tup =
  let (tup', i) = uncons tup
  in ((,) $! tupRepIdx uncons n tup') $! fromIntegral @Int @Int64 i

ixUncons :: Index (S n) -> (Index n, Int)
ixUncons (IxCons idx i) = (idx, i)

shUncons :: Shape (S n) -> (Shape n, Int)
shUncons (ShCons idx i) = (idx, i)