path: root/src/Interpreter.hs

{-# 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)
import Data.Bifunctor (bimap)
import Data.Char (isSpace)
import Data.Functor.Identity
import Data.Kind (Type)
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 trunc s | length s > lenlimit = take (lenlimit - 3) s ++ "..."
              | otherwise           = 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 _ e1 e2 -> do
    initval <- interpret' env e1
    withAccum (typeOf e1) (typeOf e2) initval $ \accum ->
      interpret' (Value accum `SCons` env) e2
  EAccum _ i e1 e2 e3 -> do
    let STAccum t = typeOf e3
    idx <- interpret' env e1
    val <- interpret' env e2
    accum <- interpret' env e3
    accumAddSparse t i 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 i a b -> do
    a' <- interpret' env a
    b' <- interpret' env b
    return $ onehotD2 i 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 :: SNat i -> STy t -> Rep (AcIdx (D2 t) i) -> Rep (AcVal (D2 t) i) -> Rep (D2 t)
onehotD2 SZ _ () v = v
onehotD2 _ STNil _ _ = ()
onehotD2 (SS SZ    ) (STPair _  _ ) ()          val         = Just val
onehotD2 (SS (SS i)) (STPair t1 t2) (Left  idx) (Left  val) = Just (onehotD2 i t1 idx val, zeroD2 t2)
onehotD2 (SS (SS i)) (STPair t1 t2) (Right idx) (Right val) = Just (zeroD2 t1, onehotD2 i t2 idx val)
onehotD2 (SS _     ) (STPair _  _ ) _           _           = error "onehotD2: pair: mismatched index and value"
onehotD2 (SS SZ    ) (STEither _  _ ) ()          val         = Just val
onehotD2 (SS (SS i)) (STEither t1 _ ) (Left  idx) (Left  val) = Just (Left (onehotD2 i t1 idx val))
onehotD2 (SS (SS i)) (STEither _  t2) (Right idx) (Right val) = Just (Right (onehotD2 i t2 idx val))
onehotD2 (SS _     ) (STEither _  _ ) _           _           = error "onehotD2: either: mismatched index and value"
onehotD2 (SS i     ) (STMaybe t) idx val = Just (onehotD2 i t idx val)
onehotD2 (SS i     ) (STArr n t) idx val = runIdentity $
  onehotArray (d2 t) (\i' idx' v' -> Identity (onehotD2 i' t idx' v')) (Identity (zeroD2 t)) n (SS i) idx val
onehotD2 SS{} STScal{} _ _ = error "onehotD2: cannot index into scalar"
onehotD2 _ STAccum{} _ _ = error "onehotD2: cannot index into accumulator"

withAccum :: STy t -> STy a -> Rep t -> (RepAcSparse t -> AcM s (Rep a)) -> AcM s (Rep a, Rep t)
withAccum t _ initval f = AcM $ do
  accum <- newAcSparse t SZ () initval
  out <- case f accum of AcM m -> m
  val <- readAcSparse t accum
  return (out, val)

newAcZero :: STy t -> IO (RepAcSparse t)
newAcZero = \case
  STNil -> return ()
  STPair t1 t2 -> newIORef =<< (,) <$> newAcZero t1 <*> newAcZero t2
  STMaybe _ -> newIORef Nothing
  STArr n _ -> newIORef (emptyArray n)
  STScal sty -> case sty of
    STI32 -> newIORef 0
    STI64 -> newIORef 0
    STF32 -> newIORef 0.0
    STF64 -> newIORef 0.0
    STBool -> error "Accumulator of Bool"
  STAccum{} -> error "Nested accumulators"
  STEither{} -> error "Bare Either in accumulator"

-- | Inverted index: the outermost index is at the /outside/ of this list.
data PartialInvIndex n m where
  PIIxEnd :: PartialInvIndex m m
  PIIxCons :: Int -> PartialInvIndex n m -> PartialInvIndex (S n) m

-- | Inverted shapey thing: the outermost dimension is at the /outside/ of this list.
data Inverted (f :: Nat -> Type) n where
  InvNil :: Inverted f Z
  InvCons :: Int -> Inverted f n -> Inverted f (S n)

type InvShape = Inverted Shape
type InvIndex = Inverted Index

class Shapey f where
  shapeyNil :: f Z
  shapeyCons :: f n -> Int -> f (S n)
  shapeyCase :: f n -> (n ~ Z => r) -> (forall m. n ~ S m => f m -> Int -> r) -> r
instance Shapey Index where
  shapeyNil = IxNil
  shapeyCons = IxCons
  shapeyCase IxNil k0 _ = k0
  shapeyCase (IxCons idx i) _ k1 = k1 idx i
instance Shapey Shape where
  shapeyNil = ShNil
  shapeyCons = ShCons
  shapeyCase ShNil k0 _ = k0
  shapeyCase (ShCons sh n) _ k1 = k1 sh n

invert :: forall f n. Shapey f => f n -> Inverted f n
invert | Refl <- lemPlusZero @n = flip go InvNil
  where
    go :: forall n' m. f n' -> Inverted f m -> Inverted f (n' + m)
    go sh ish = shapeyCase sh
                  ish
                  (\sh' n -> case lemPlusSuccRight @n' @m of Refl -> go sh' (InvCons n ish))

uninvert :: forall f n. Shapey f => Inverted f n -> f n
uninvert = go shapeyNil
  where
    go :: forall n' m. f n' -> Inverted f m -> f (n' + m)
    go sh InvNil | Refl <- lemPlusZero @n' = sh
    go sh (InvCons n (ish :: Inverted f predm)) | Refl <- lemPlusSuccRight @n' @predm = go (shapeyCons sh n) ish

piindexMatch :: PartialInvIndex n m -> InvIndex n -> Maybe (InvIndex m)
piindexMatch PIIxEnd ix = Just ix
piindexMatch (PIIxCons i pix) (InvCons i' ix)
  | i == i' = piindexMatch pix ix
  | otherwise = Nothing

piindexConcat :: PartialInvIndex n m -> InvIndex m -> InvIndex n
piindexConcat PIIxEnd ix = ix
piindexConcat (PIIxCons i pix) ix = InvCons i (piindexConcat pix ix)

newAcSparse :: STy t -> SNat i -> Rep (AcIdx t i) -> Rep (AcVal t i) -> IO (RepAcSparse t)
newAcSparse typ SZ () val = case typ of
  STNil -> return ()
  STPair t1 t2 -> newIORef =<< (,) <$> newAcSparse t1 SZ () (fst val) <*> newAcSparse t2 SZ () (snd val)
  STMaybe t -> newIORef =<< traverse (newAcDense t SZ ()) val
  STArr _ t -> newIORef =<< traverse (newAcSparse t SZ ()) val
  STScal{} -> newIORef val
  STAccum{} -> error "Nested accumulators"
  STEither{} -> error "Bare Either in accumulator"
newAcSparse typ (SS dep) idx val = case typ of
  STNil -> return ()
  STPair t1 t2 -> newIORef =<< case (idx, val) of
                                 (Left idx', Left val') -> (,) <$> newAcSparse t1 dep idx' val' <*> newAcZero t2
                                 (Right idx', Right val') -> (,) <$> newAcZero t1 <*> newAcSparse t2 dep idx' val'
                                 _ -> error "Index/value mismatch in newAc pair"
  STMaybe t -> newIORef =<< Just <$> newAcDense t dep idx val
  STArr dim (t :: STy t) -> newIORef =<< newAcArray dim t (SS dep) idx val
  STScal{} -> error "Cannot index into scalar"
  STAccum{} -> error "Nested accumulators"
  STEither{} -> error "Bare Either in accumulator"

newAcArray :: SNat n -> STy t -> SNat i -> Rep (AcIdx (TArr n t) i) -> Rep (AcVal (TArr n t) i) -> IO (Array n (RepAcSparse t))
newAcArray n t = onehotArray t (newAcSparse t) (newAcZero t) n

onehotArray :: Monad m
            => STy t
            -> (forall n'. SNat n' -> Rep (AcIdx t n') -> Rep (AcVal t n') -> m v)  -- ^ the "one"
            -> m v  -- ^ generate a zero value for elsewhere
            -> SNat n -> SNat i -> Rep (AcIdx (TArr n t) i) -> Rep (AcVal (TArr n t) i) -> m (Array n v)
onehotArray _ mkone _ _ SZ _ val =
  traverse (mkone SZ ()) val
onehotArray (_ :: STy t) mkone mkzero dim dep@SS{} idx val = do
  let sh = unTupRepIdx ShNil ShCons dim (fst val)
  go mkone dep dim idx (snd val) $ \arr position ->
    arrayGenerateM sh (\i -> case uninvert <$> piindexMatch position (invert i) of
                               Just i' -> return $ arr `arrayIndex` i'
                               Nothing -> mkzero)
  where
    go :: Monad m
       => (forall n'. SNat n' -> Rep (AcIdx t n') -> Rep (AcVal t n') -> m v)
       -> SNat i -> SNat n -> Rep (AcIdx (TArr n t) i) -> Rep (AcValArr n t i)
       -> (forall n'. Array n' v -> PartialInvIndex n n' -> m r) -> m r
    go mk SZ _ () val' k = arrayMapM (mk SZ ()) val' >>= \arr -> k arr PIIxEnd
    go mk (SS dep') SZ idx' val' k = mk dep' idx' val' >>= \arr -> k (arrayUnit arr) PIIxEnd
    go mk (SS dep') (SS dim') (i, idx') val' k =
      go mk dep' dim' idx' val' $ \arr pish ->
        k arr (PIIxCons (fromIntegral @Int64 @Int i) pish)

newAcDense :: STy t -> SNat i -> Rep (AcIdx t i) -> Rep (AcVal t i) -> IO (RepAcDense t)
newAcDense typ SZ () val = case typ of
  STPair t1 t2 -> (,) <$> newAcSparse t1 SZ () (fst val) <*> newAcSparse t2 SZ () (snd val)
  STEither t1 t2 -> case val of
    Left x -> Left <$> newAcSparse t1 SZ () x
    Right y -> Right <$> newAcSparse t2 SZ () y
  _ -> error "newAcDense: invalid dense type"
newAcDense typ (SS dep) idx val = case typ of
  STPair t1 t2 ->
    case (idx, val) of
      (Left idx', Left val') -> (,) <$> newAcSparse t1 dep idx' val' <*> newAcZero t2
      (Right idx', Right val') -> (,) <$> newAcZero t1 <*> newAcSparse t2 dep idx' val'
      _ -> error "Index/value mismatch in newAc pair"
  STEither t1 t2 ->
    case (idx, val) of
      (Left idx', Left val') -> Left <$> newAcSparse t1 dep idx' val'
      (Right idx', Right val') -> Right <$> newAcSparse t2 dep idx' val'
      _ -> error "Index/value mismatch in newAc either"
  _ -> error "newAcDense: invalid dense type"

readAcSparse :: STy t -> RepAcSparse t -> IO (Rep t)
readAcSparse typ val = case typ of
  STNil -> return ()
  STPair t1 t2 -> do
    (a, b) <- readIORef val
    (,) <$> readAcSparse t1 a <*> readAcSparse t2 b
  STMaybe t -> traverse (readAcDense t) =<< readIORef val
  STArr _ t -> traverse (readAcSparse t) =<< readIORef val
  STScal{} -> readIORef val
  STAccum{} -> error "Nested accumulators"
  STEither{} -> error "Bare Either in accumulator"

readAcDense :: STy t -> RepAcDense t -> IO (Rep t)
readAcDense typ val = case typ of
  STPair t1 t2 -> (,) <$> readAcSparse t1 (fst val) <*> readAcSparse t2 (snd val)
  STEither t1 t2 -> case val of
    Left x -> Left <$> readAcSparse t1 x
    Right y -> Right <$> readAcSparse t2 y
  _ -> error "readAcDense: invalid dense type"

accumAddSparse :: STy t -> SNat i -> RepAcSparse t -> Rep (AcIdx t i) -> Rep (AcVal t i) -> AcM s ()
accumAddSparse typ SZ ref () val = case typ of
  STNil -> return ()
  STPair t1 t2 -> AcM $ do
    (r1, r2) <- readIORef ref
    unAcM $ accumAddSparse t1 SZ r1 () (fst val)
    unAcM $ accumAddSparse t2 SZ r2 () (snd val)
  STMaybe t ->
    case val of
      Nothing -> return ()
      Just val' ->
        -- Try adding val' to 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 newac <- newAcDense t SZ () val'
                                          join $ atomicModifyIORef' ref $ \ac2 -> case ac2 of
                                                   Nothing -> (Just newac, return ())
                                                   Just ac2' -> bimap Just unAcM (accumAddDense t SZ ac2' () val'))
                       -- Yep, ref already had a value in there, so we can just add
                       -- val' to it recursively.
                       Just ac' -> bimap Just unAcM (accumAddDense t SZ ac' () val')
  STArr _ t -> AcM $ do
    refs <- readIORef ref
    case (shapeSize (arrayShape refs), shapeSize (arrayShape val)) of
      (_, 0) -> return ()
      (0, _) -> do
        newrefarr <- traverse (newAcSparse t SZ ()) val
        join $ atomicModifyIORef' ref $ \refarr ->
                 if shapeSize (arrayShape refarr) == 0
                   then (newrefarr, return ())
                   else -- someone was faster than us in initialising the reference!
                        (refarr, unAcM $ accumAddSparse typ SZ ref () val)  -- just try again from the start (dropping newrefarr for the GC to clean up)
      _ | arrayShape refs == arrayShape val ->
            sequence_ [unAcM $ accumAddSparse t SZ (arrayIndexLinear refs i) () (arrayIndexLinear val i)
                      | i <- [0 .. shapeSize (arrayShape val) - 1]]
        | otherwise -> error "Array shape mismatch in accum add"
  STScal sty -> AcM $ case sty of
    STI32 -> atomicModifyIORef' ref (\x -> (x + val, ()))
    STI64 -> atomicModifyIORef' ref (\x -> (x + val, ()))
    STF32 -> atomicModifyIORef' ref (\x -> (x + val, ()))
    STF64 -> atomicModifyIORef' ref (\x -> (x + val, ()))
    STBool -> error "Accumulator of Bool"
  STAccum{} -> error "Nested accumulators"
  STEither{} -> error "Bare Either in accumulator"

accumAddSparse typ (SS dep) ref idx val = case typ of
  STNil -> return ()
  STPair t1 t2 -> AcM $ do
    (ref1, ref2) <- readIORef ref
    case (idx, val) of
      (Left idx', Left val') -> unAcM $ accumAddSparse t1 dep ref1 idx' val'
      (Right idx', Right val') -> unAcM $ accumAddSparse t2 dep ref2 idx' val'
      _ -> error "Index/value mismatch in pair accumulator add"
  STMaybe t ->
    AcM $ join $ atomicModifyIORef' ref $ \case
                   -- Oops, ref's contents was still sparse. Have to initialise
                   -- it first, then try again.
                   Nothing -> (Nothing, do newac <- newAcDense t dep idx val
                                           join $ atomicModifyIORef' ref $ \ac2 -> case ac2 of
                                                    Nothing -> (Just newac, return ())
                                                    Just ac2' -> bimap Just unAcM (accumAddDense t dep ac2' idx val))
                   -- Yep, ref already had a value in there, so we can just add
                   -- val' to it recursively.
                   Just ac -> bimap Just unAcM (accumAddDense t dep ac idx val)
  STArr dim (t :: STy t) -> AcM $ do
    refs <- readIORef ref
    if shapeSize (arrayShape refs) == 0
      then do newrefarr <- newAcArray dim t (SS dep) idx val
              join $ atomicModifyIORef' ref $ \refarr ->
                       if shapeSize (arrayShape refarr) == 0
                         then (newrefarr, return ())
                         else -- someone was faster than us in initialising the reference!
                              (refarr, unAcM $ accumAddSparse typ (SS dep) ref idx val)  -- just try again from the start (dropping newrefarr for the GC to clean up)
      else do let sh = unTupRepIdx ShNil ShCons dim (fst val)
              go (SS dep) (invert sh) idx (snd val)
                (\j index idxj valj -> unAcM $ accumAddSparse t j (refs `arrayIndex` index) idxj valj)
                (\piix subsh val' -> unAcM $ sequence_
                                       [accumAddSparse t SZ (refs `arrayIndex` uninvert (piindexConcat piix (invert subix)))
                                                       () (val' `arrayIndex` subix)
                                       | subix <- enumShape subsh])
    where
      go :: SNat i -> InvShape n -> Rep (AcIdx (TArr n t) i) -> Rep (AcValArr n t i)
         -> (forall j. SNat j -> Index n -> Rep (AcIdx t j) -> Rep (AcVal t j) -> r)  -- ^ Indexing into element of the array
         -> (forall m. PartialInvIndex n m -> Shape m -> Rep (TArr m t) -> r)  -- ^ Accumulating onto a subarray
         -> r
      go SZ        ish             ()        val' _  k0 = k0 PIIxEnd (uninvert ish) val'  -- ^ Ran out of AcIdx: accumulating onto subarray
      go (SS dep') InvNil          idx'      val' kj _  = kj dep' IxNil idx' val'  -- ^ Ran out of array dimensions: accumulating into (part of) element
      go (SS dep') (InvCons _ ish) (i, idx') val' kj k0 =
        go dep' ish idx' val'
          (\j index idxj valj -> kj j (IxCons index (fromIntegral @Int64 @Int i)) idxj valj)
          (\pidxm shm valm -> k0 (PIIxCons (fromIntegral @Int64 @Int i) pidxm) shm valm)
  STScal{} -> error "Cannot index into scalar"
  STAccum{} -> error "Nested accumulators"
  STEither{} -> error "Bare Either in accumulator"

accumAddDense :: forall t i s. STy t -> SNat i -> RepAcDense t -> Rep (AcIdx t i) -> Rep (AcVal t i) -> (RepAcDense t, AcM s ())
accumAddDense typ SZ ref () val = case typ of
  STPair t1 t2 ->
    (ref, do accumAddSparse t1 SZ (fst ref) () (fst val)
             accumAddSparse t2 SZ (snd ref) () (snd val))
  STEither t1 t2 ->
    case (ref, val) of
      (Left ref', Left val') -> (ref, accumAddSparse t1 SZ ref' () val')
      (Right ref', Right val') -> (ref, accumAddSparse t2 SZ ref' () val')
      _ -> error "Mismatched Either in accumAddDense either"
  _ -> error "accumAddDense: invalid dense type"

accumAddDense typ (SS dep) ref idx val = case typ of
  STPair t1 t2 ->
    case (idx, val) of
      (Left idx', Left val') -> (ref, accumAddSparse t1 dep (fst ref) idx' val')
      (Right idx', Right val') -> (ref, accumAddSparse t2 dep (snd ref) idx' val')
      _ -> error "Mismatched Either in accumAddDense pair"
  STEither t1 t2 ->
    case (ref, idx, val) of
      (Left ref', Left idx', Left val') -> (Left ref', accumAddSparse t1 dep ref' idx' val')
      (Right ref', Right idx', Right val') -> (Right ref', accumAddSparse t2 dep ref' idx' val')
      _ -> error "Mismatched Either in accumAddDense either"
  _ -> error "accumAddDense: invalid dense type"


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)