{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeApplications #-}
module Compile (compile, debugCSource, debugRefc, emitChecks) where

import Control.Applicative (empty)
import Control.Monad (forM_, when, replicateM)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Maybe
import Control.Monad.Trans.State.Strict
import Control.Monad.Trans.Writer.CPS
import Data.Bifunctor (first)
import Data.Char (ord)
import Data.Foldable (toList)
import Data.Functor.Const
import qualified Data.Functor.Product as Product
import Data.Functor.Product (Product)
import Data.List (foldl1', intersperse, intercalate)
import qualified Data.Map.Strict as Map
import Data.Maybe (fromMaybe)
import qualified Data.Set as Set
import Data.Set (Set)
import Data.Some
import qualified Data.Vector as V
import Foreign
import GHC.Exts (int2Word#, addr2Int#)
import GHC.Num (integerFromWord#)
import GHC.Ptr (Ptr(..))
import Numeric (showHex)
import System.IO (hPutStrLn, stderr)

import Prelude hiding ((^))
import qualified Prelude

import Array
import AST
import AST.Pretty (ppSTy, ppExpr)
import qualified CHAD.Types as CHAD
import Compile.Exec
import Data
import Interpreter.Rep
import qualified Util.IdGen as IdGen


-- In shape and index arrays, the innermost dimension is on the right (last index).

-- TODO: test that I'm properly incrementing and decrementing refcounts in all required places


-- | Print the compiled AST
debugPrintAST :: Bool; debugPrintAST = toEnum 0
-- | Print the generated C source
debugCSource :: Bool; debugCSource = toEnum 0
-- | Print extra stuff about reference counts of arrays
debugRefc :: Bool; debugRefc = toEnum 0
-- | Print some shape-related information
debugShapes :: Bool; debugShapes = toEnum 0
-- | Print information on allocation
debugAllocs :: Bool; debugAllocs = toEnum 0
-- | Emit extra C code that checks stuff
emitChecks :: Bool; emitChecks = toEnum 0

compile :: SList STy env -> Ex env t
        -> IO (SList Value env -> IO (Rep t))
compile = \env expr -> do
  let source = compileToString env expr
  when debugPrintAST $ hPutStrLn stderr $ "Compiled AST: <<<\n" ++ ppExpr env expr ++ "\n>>>"
  when debugCSource $ hPutStrLn stderr $ "Generated C source: <<<\n\x1B[2m" ++ source ++ "\x1B[0m>>>"
  lib <- buildKernel source ["kernel"]

  let arg_metrics = reverse (unSList metricsSTy env)
      (arg_offsets, result_offset) = computeStructOffsets arg_metrics
      result_type = typeOf expr
      result_size = sizeofSTy result_type

  return $ \val -> do
    allocaBytes (result_offset + result_size) $ \ptr -> do
      let args = zip (reverse (unSList Some (slistZip env val))) arg_offsets
      serialiseArguments args ptr $ do
        callKernelFun "kernel" lib ptr
        deserialise result_type ptr result_offset
  where
    serialiseArguments :: [(Some (Product STy Value), Int)] -> Ptr () -> IO r -> IO r
    serialiseArguments ((Some (Product.Pair t (Value arg)), off) : args) ptr k =
      serialise t arg ptr off $
        serialiseArguments args ptr k
    serialiseArguments _ _ k = k


data StructDecl = StructDecl
    String  -- ^ name
    String  -- ^ contents
    String  -- ^ comment
  deriving (Show)

data Stmt
  = SVarDecl Bool String String CExpr  -- ^ const, type, variable name, right-hand side
  | SVarDeclUninit String String  -- ^ type, variable name (no initialiser)
  | SAsg String CExpr  -- ^ variable name, right-hand side
  | SBlock (Bag Stmt)
  | SIf CExpr (Bag Stmt) (Bag Stmt)
  | SLoop String String CExpr CExpr (Bag Stmt)  -- ^ for (<type> <name> = <expr>; name < <expr>; name++) {<stmts>}
  | SVerbatim String  -- ^ no implicit ';', just printed as-is
  deriving (Show)

data CExpr
  = CELit String  -- ^ inserted as-is, assumed no parentheses needed
  | CEStruct String [(String, CExpr)]  -- ^ struct construction literal: `(name){.field=expr}`
  | CEProj CExpr String  -- ^ field projection: expr.field
  | CEPtrProj CExpr String  -- ^ field projection through pointer: expr->field
  | CEAddrOf CExpr  -- ^ &expr
  | CEIndex CExpr CExpr  -- ^ expr[expr]
  | CECall String [CExpr]  -- ^ function(arg1, ..., argn)
  | CEBinop CExpr String CExpr  -- ^ expr + expr
  | CEIf CExpr CExpr CExpr  -- ^ expr ? expr : expr
  | CECast String CExpr  -- ^ (<type)<expr>
  deriving (Show)

printStructDecl :: StructDecl -> ShowS
printStructDecl (StructDecl name contents comment) =
  showString "typedef struct { " . showString contents . showString " } " . showString name
  . showString ";" . (if null comment then id else showString ("  // " ++ comment))

printStmt :: Int -> Stmt -> ShowS
printStmt indent = \case
  SVarDecl cnst typ name rhs -> showString (typ ++ " " ++ (if cnst then "const " else "") ++ name ++ " = ") . printCExpr 0 rhs . showString ";"
  SVarDeclUninit typ name -> showString (typ ++ " " ++ name ++ ";")
  SAsg name rhs -> showString (name ++ " = ") . printCExpr 0 rhs . showString ";"
  SBlock stmts
    | null stmts -> showString "{}"
    | otherwise ->
        showString "{"
        . compose [showString ("\n" ++ replicate (2*indent+2) ' ') . printStmt (indent+1) stmt | stmt <- toList stmts]
        . showString ("\n" ++ replicate (2*indent) ' ' ++ "}")
  SIf cond b1 b2 ->
    showString "if (" . printCExpr 0 cond . showString ") "
    . printStmt indent (SBlock b1)
    . (if null b2 then id else showString " else " . printStmt indent (SBlock b2))
  SLoop typ name e1 e2 stmts ->
    showString ("for (" ++ typ ++ " " ++ name ++ " = ")
    . printCExpr 0 e1 . showString ("; " ++ name ++ " < ") . printCExpr 6 e2
    . showString ("; " ++ name ++ "++) ")
    . printStmt indent (SBlock stmts)
  SVerbatim s -> showString s

-- d values:
-- * 0: top level
-- * 1: in 1st or 2nd component of a ternary operator (technically same as top level, but readability)
-- * 2-...: various operators (see precTable)
-- * 80: address-of operator (&)
-- * 98: inside unknown operator
-- * 99: left of a field projection
-- Unlisted operators are conservatively written with full parentheses.
printCExpr :: Int -> CExpr -> ShowS
printCExpr d = \case
  CELit s -> showString s
  CEStruct name pairs ->
    showParen (d >= 99) $
      showString ("(" ++ name ++ "){")
      . compose (intersperse (showString ", ") [showString ("." ++ n ++ " = ") . printCExpr 0 e
                                               | (n, e) <- pairs])
      . showString "}"
  CEProj e name -> printCExpr 99 e . showString ("." ++ name)
  CEPtrProj e name -> printCExpr 99 e . showString ("->" ++ name)
  CEAddrOf e -> showParen (d > 80) $ showString "&" . printCExpr 80 e
  CEIndex e1 e2 -> printCExpr 99 e1 . showString "[" . printCExpr 0 e2 . showString "]"
  CECall n es ->
    showString (n ++ "(") . compose (intersperse (showString ", ") (map (printCExpr 0) es)) . showString ")"
  CEBinop e1 n e2 ->
    let mprec = Map.lookup n precTable
        p = maybe (-1) fst mprec  -- precedence of this operator
        (d1, d2) = maybe (98, 98) snd mprec  -- precedences for the arguments
    in showParen (d > p) $
         printCExpr d1 e1 . showString (" " ++ n ++ " ") . printCExpr d2 e2
  CEIf e1 e2 e3 ->
    showParen (d > 0) $
      printCExpr 1 e1 . showString " ? " . printCExpr 1 e2 . showString " : " . printCExpr 0 e3
  CECast typ e ->
    showParen (d > 98) $ showString ("(" ++ typ ++ ")") . printCExpr 98 e
  where
    precTable = Map.fromList
      [("||", (2, (2, 2)))
      ,("&&", (3, (3, 3)))
      ,("==", (4, (5, 5)))
      ,("!=", (4, (5, 5)))
      ,("<", (5, (6, 6)))  -- Note: this precedence is used in the printing of SLoop
      ,(">", (5, (6, 6)))
      ,("<=", (5, (6, 6)))
      ,(">=", (5, (6, 6)))
      ,("+", (6, (6, 7)))
      ,("-", (6, (6, 7)))
      ,("*", (7, (7, 8)))
      ,("/", (7, (7, 8)))
      ,("%", (7, (7, 8)))]

repSTy :: STy t -> String
repSTy (STScal st) = case st of
  STI32 -> "int32_t"
  STI64 -> "int64_t"
  STF32 -> "float"
  STF64 -> "double"
  STBool -> "uint8_t"
repSTy t = genStructName t

genStructName :: STy t -> String
genStructName = \t -> "ty_" ++ gen t where
  -- all tags start with a letter, so the array mangling is unambiguous.
  gen :: STy t -> String
  gen STNil = "n"
  gen (STPair a b) = 'P' : gen a ++ gen b
  gen (STEither a b) = 'E' : gen a ++ gen b
  gen (STMaybe t) = 'M' : gen t
  gen (STArr n t) = "A" ++ show (fromSNat n) ++ gen t
  gen (STScal st) = case st of
    STI32 -> "i"
    STI64 -> "j"
    STF32 -> "f"
    STF64 -> "d"
    STBool -> "b"
  gen (STAccum t) = 'C' : gen t

-- | This function generates the actual struct declarations for each of the
-- types in our language. It thus implicitly "documents" the layout of the
-- types in the C translation.
genStruct :: String -> STy t -> [StructDecl]
genStruct name topty = case topty of
  STNil ->
    [StructDecl name "" com]
  STPair a b ->
    [StructDecl name (repSTy a ++ " a; " ++ repSTy b ++ " b;") com]
  STEither a b ->  -- 0 -> l, 1 -> r
    [StructDecl name ("uint8_t tag; union { " ++ repSTy a ++ " l; " ++ repSTy b ++ " r; };") com]
  STMaybe t ->  -- 0 -> nothing, 1 -> just
    [StructDecl name ("uint8_t tag; " ++ repSTy t ++ " j;") com]
  STArr n t ->
    -- The buffer is trailed by a VLA for the actual array data.
    [StructDecl (name ++ "_buf") ("size_t sh[" ++ show (fromSNat n) ++ "]; size_t refc; " ++ repSTy t ++ " xs[];") ""
    ,StructDecl name (name ++ "_buf *buf;") com]
  STScal _ ->
    []
  STAccum t ->
    [StructDecl (name ++ "_buf") (repSTy (CHAD.d2 t) ++ " ac;") ""
    ,StructDecl name (name ++ "_buf *buf;") com]
  where
    com = ppSTy 0 topty

-- State: already-generated (skippable) struct names
-- Writer: the structs in declaration order
genStructs :: STy t -> WriterT (Bag StructDecl) (State (Set String)) ()
genStructs ty = do
  let name = genStructName ty
  seen <- lift $ gets (name `Set.member`)

  if seen
    then pure ()
    else do
      -- already mark this struct as generated now, so we don't generate it
      -- twice (unnecessary because no recursive types, but y'know)
      lift $ modify (Set.insert name)

      () <- case ty of
        STNil -> pure ()
        STPair a b -> genStructs a >> genStructs b
        STEither a b -> genStructs a >> genStructs b
        STMaybe t -> genStructs t
        STArr _ t -> genStructs t
        STScal _ -> pure ()
        STAccum t -> genStructs (CHAD.d2 t)

      tell (BList (genStruct name ty))

genAllStructs :: Foldable t => t Ty -> [StructDecl]
genAllStructs tys = toList $ evalState (execWriterT (mapM_ (\t -> case reSTy t of Some t' -> genStructs t') tys)) mempty

data CompState = CompState
  { csStructs :: Set Ty
  , csTopLevelDecls :: Bag String
  , csStmts :: Bag Stmt
  , csNextId :: Int }
  deriving (Show)

newtype CompM a = CompM (State CompState a)
  deriving newtype (Functor, Applicative, Monad)

runCompM :: CompM a -> (a, CompState)
runCompM (CompM m) = runState m (CompState mempty mempty mempty 1)

class Monad m => MonadNameGen m where genId :: m Int
instance MonadNameGen CompM where genId = CompM $ state $ \s -> (csNextId s, s { csNextId = csNextId s + 1 })
instance MonadNameGen IdGen.IdGen where genId = IdGen.genId
instance MonadNameGen m => MonadNameGen (MaybeT m) where genId = MaybeT (Just <$> genId)

genName' :: MonadNameGen m => String -> m String
genName' "" = genName
genName' prefix = (prefix ++) . show <$> genId

genName :: MonadNameGen m => m String
genName = genName' "x"

onlyIdGen :: IdGen.IdGen a -> CompM a
onlyIdGen m = CompM $ do
  i1 <- gets csNextId
  let (res, i2) = IdGen.runIdGen' i1 m
  modify (\s -> s { csNextId = i2 })
  return res

emit :: Stmt -> CompM ()
emit stmt = CompM $ modify $ \s -> s { csStmts = csStmts s <> pure stmt }

scope :: CompM a -> CompM (a, Bag Stmt)
scope m = do
  stmts <- CompM $ state $ \s -> (csStmts s, s { csStmts = mempty })
  res <- m
  innerStmts <- CompM $ state $ \s -> (csStmts s, s { csStmts = stmts })
  return (res, innerStmts)

emitStruct :: STy t -> CompM String
emitStruct ty = CompM $ do
  modify $ \s -> s { csStructs = Set.insert (unSTy ty) (csStructs s) }
  return (genStructName ty)

emitTLD :: String -> CompM ()
emitTLD decl = CompM $ modify $ \s -> s { csTopLevelDecls = csTopLevelDecls s <> pure decl }

nameEnv :: SList f env -> SList (Const String) env
nameEnv = flip evalState (0::Int) . slistMapA (\_ -> state $ \i -> (Const ("arg" ++ show i), i + 1))

compileToString :: SList STy env -> Ex env t -> String
compileToString env expr =
  let args = nameEnv env
      (res, s) = runCompM (compile' args expr)
      structs = genAllStructs (csStructs s <> Set.fromList (unSList unSTy env))

      (arg_pairs, arg_metrics) =
        unzip $ reverse (unSList (\(Product.Pair t (Const n)) -> ((n, repSTy t), metricsSTy t))
                                 (slistZip env args))
      (arg_offsets, result_offset') = computeStructOffsets arg_metrics
      result_offset = align (alignmentSTy (typeOf expr)) result_offset'
  in ($ "") $ compose
       [showString "#include <stdio.h>\n"
       ,showString "#include <stdint.h>\n"
       ,showString "#include <inttypes.h>\n"
       ,showString "#include <stdlib.h>\n"
       ,showString "#include <string.h>\n"
       ,showString "#include <math.h>\n\n"
       ,compose [printStructDecl sd . showString "\n" | sd <- structs]
       ,showString "\n"
       ,showString "static void* malloc_instr(size_t n) {\n"
       ,showString "  void *ptr = malloc(n);\n"
       ,if debugAllocs then showString "printf(\"[chad-kernel] malloc(%zu) -> %p\\n\", n, ptr);\n"
                       else id
       ,showString "  return ptr;\n"
       ,showString "}\n"
       ,showString "static void* calloc_instr(size_t n) {\n"
       ,showString "  void *ptr = calloc(n, 1);\n"
       ,if debugAllocs then showString "printf(\"[chad-kernel] calloc(%zu) -> %p\\n\", n, ptr);\n"
                       else id
       ,showString "  return ptr;\n"
       ,showString "}\n"
       ,showString "static void free_instr(void *ptr) {\n"
       ,if debugAllocs then showString "printf(\"[chad-kernel] free(%p)\\n\", ptr);\n"
                       else id
       ,showString "  free(ptr);\n"
       ,showString "}\n\n"
       ,compose [showString str . showString "\n\n" | str <- toList (csTopLevelDecls s)]
       ,showString $
          "static " ++ repSTy (typeOf expr) ++ " typed_kernel(" ++
            intercalate ", " (reverse (unSList (\(Product.Pair t (Const n)) -> repSTy t ++ " " ++ n) (slistZip env args))) ++
            ") {\n"
       ,compose [showString "  " . printStmt 1 st . showString "\n" | st <- toList (csStmts s)]
       ,showString "  return " . printCExpr 0 res . showString ";\n}\n\n"
       ,showString "void kernel(void *data) {\n"
        -- Some code here assumes that we're on a 64-bit system, so let's check that
       ,showString "  if (sizeof(void*) != 8 || sizeof(size_t) != 8) { fprintf(stderr, \"Only 64-bit systems supported\\n\"); abort(); }\n"
       ,if debugRefc then showString "  fprintf(stderr, \"[chad-kernel] Start\\n\");\n"
                     else id
       ,showString $ "  *(" ++ repSTy (typeOf expr) ++ "*)(data + " ++ show result_offset ++ ") = typed_kernel(" ++
                       concat (map (\((arg, typ), off, idx) ->
                                       "\n    *(" ++ typ ++ "*)(data + " ++ show off ++ ")"
                                       ++ (if idx < length arg_pairs - 1 then "," else "")
                                       ++ "  // " ++ arg)
                                   (zip3 arg_pairs arg_offsets [0::Int ..])) ++
                       "\n  );\n"
       ,if debugRefc then showString "  fprintf(stderr, \"[chad-kernel] Return\\n\");\n"
                     else id
       ,showString "}\n"]

-- | Takes list of metrics (alignment, sizeof).
-- Returns (offsets, size of struct).
computeStructOffsets :: [(Int, Int)] -> ([Int], Int)
computeStructOffsets = go 0 0
  where
    go off maxal [(al, sz)] =
      ([off], align (max maxal al) (off + sz))
    go off maxal ((al, sz) : pairs@((al2,_):_)) =
      first (off :) $ go (align al2 (off + sz)) (max maxal al) pairs
    go _   _     [] = ([], 0)

-- | Assumes that this is called at the correct alignment.
serialise :: STy t -> Rep t -> Ptr () -> Int -> IO r -> IO r
serialise topty topval ptr off k =
  -- TODO: this code is quadratic in the depth of the type because of the alignment/sizeOf calls
  case (topty, topval) of
    (STNil, ()) -> k
    (STPair a b, (x, y)) ->
      serialise a x ptr off $
        serialise b y ptr (align (alignmentSTy b) (off + sizeofSTy a)) k
    (STEither a _, Left x) -> do
      pokeByteOff ptr off (0 :: Word8)  -- alignment of (a + b) is alignment of (union {a b})
      serialise a x ptr (off + alignmentSTy topty) k
    (STEither _ b, Right y) -> do
      pokeByteOff ptr off (1 :: Word8)
      serialise b y ptr (off + alignmentSTy topty) k
    (STMaybe _, Nothing) -> do
      pokeByteOff ptr off (0 :: Word8)
      k
    (STMaybe t, Just x) -> do
      pokeByteOff ptr off (1 :: Word8)
      serialise t x ptr (off + alignmentSTy t) k
    (STArr n t, Array sh vec) -> do
      let eltsz = sizeofSTy t
      allocaBytes (fromSNat n * 8 + 8 + shapeSize sh * eltsz) $ \bufptr -> do
        when debugRefc $
          hPutStrLn stderr $ "[chad-serialise] Allocating input buffer " ++ showPtr bufptr
        pokeByteOff ptr off bufptr

        pokeShape bufptr 0 n sh
        pokeByteOff @Word64 bufptr (8 * fromSNat n) (2 ^ 63)

        let off1 = fromSNat n * 8 + 8
            loop i
              | i == shapeSize sh = k
              | otherwise =
                  serialise t (vec V.! i) bufptr (off1 + i * eltsz) $
                    loop (i+1)
        loop 0
    (STScal sty, x) -> case sty of
      STI32 -> pokeByteOff ptr off (x :: Int32) >> k
      STI64 -> pokeByteOff ptr off (x :: Int64) >> k
      STF32 -> pokeByteOff ptr off (x :: Float) >> k
      STF64 -> pokeByteOff ptr off (x :: Double) >> k
      STBool -> pokeByteOff ptr off (fromIntegral (fromEnum x) :: Word8) >> k
    (STAccum{}, _) -> error "Cannot serialise accumulators"

-- | Assumes that this is called at the correct alignment.
deserialise :: STy t -> Ptr () -> Int -> IO (Rep t)
deserialise topty ptr off =
  -- TODO: this code is quadratic in the depth of the type because of the alignment/sizeOf calls
  case topty of
    STNil -> return ()
    STPair a b -> do
      x <- deserialise a ptr off
      y <- deserialise b ptr (align (alignmentSTy b) (off + sizeofSTy a))
      return (x, y)
    STEither a b -> do
      tag <- peekByteOff @Word8 ptr off
      if tag == 0  -- alignment of (a + b) is alignment of (union {a b})
        then Left <$> deserialise a ptr (off + alignmentSTy topty)
        else Right <$> deserialise b ptr (off + alignmentSTy topty)
    STMaybe t -> do
      tag <- peekByteOff @Word8 ptr off
      if tag == 0
        then return Nothing
        else Just <$> deserialise t ptr (off + alignmentSTy t)
    STArr n t -> do
      bufptr <- peekByteOff @(Ptr ()) ptr off
      sh <- peekShape bufptr 0 n
      refc <- peekByteOff @Word64 bufptr (8 * fromSNat n)
      when debugRefc $
        hPutStrLn stderr $ "[chad-deserialise] Got buffer " ++ showPtr bufptr ++ " at refc=" ++ show refc
      let off1 = 8 * fromSNat n + 8
          eltsz = sizeofSTy t
      arr <- Array sh <$> V.generateM (shapeSize sh) (\i -> deserialise t bufptr (off1 + i * eltsz))
      when (refc < 2 ^ 62) $ free bufptr
      return arr
    STScal sty -> case sty of
      STI32 -> peekByteOff @Int32 ptr off
      STI64 -> peekByteOff @Int64 ptr off
      STF32 -> peekByteOff @Float ptr off
      STF64 -> peekByteOff @Double ptr off
      STBool -> toEnum . fromIntegral <$> peekByteOff @Word8 ptr off
    STAccum{} -> error "Cannot serialise accumulators"

align :: Int -> Int -> Int
align a off = (off + a - 1) `div` a * a

alignmentSTy :: STy t -> Int
alignmentSTy = fst . metricsSTy

sizeofSTy :: STy t -> Int
sizeofSTy = snd . metricsSTy

-- | Returns (alignment, sizeof)
metricsSTy :: STy t -> (Int, Int)
metricsSTy STNil = (1, 0)
metricsSTy (STPair a b) =
  let (a1, s1) = metricsSTy a
      (a2, s2) = metricsSTy b
  in (max a1 a2, align (max a1 a2) (s1 + s2))
metricsSTy (STEither a b) =
  let (a1, s1) = metricsSTy a
      (a2, s2) = metricsSTy b
  in (max a1 a2, max a1 a2 + max s1 s2)  -- the union after the tag byte is aligned
metricsSTy (STMaybe t) =
  let (a, s) = metricsSTy t
  in (a, a + s)  -- the union after the tag byte is aligned
metricsSTy (STArr _ _) = (8, 8)
metricsSTy (STScal sty) = case sty of
  STI32 -> (4, 4)
  STI64 -> (8, 8)
  STF32 -> (4, 4)
  STF64 -> (8, 8)
  STBool -> (1, 1)  -- compiled to uint8_t
metricsSTy (STAccum t) = metricsSTy t

pokeShape :: Ptr () -> Int -> SNat n -> Shape n -> IO ()
pokeShape ptr off = go . fromSNat
  where
    go :: Int -> Shape n -> IO ()
    go rank = \case
      ShNil -> return ()
      sh `ShCons` n -> do
        pokeByteOff ptr (off + (rank - 1) * 8) (fromIntegral n :: Int64)
        go (rank - 1) sh

peekShape :: Ptr () -> Int -> SNat n -> IO (Shape n)
peekShape ptr off = \case
  SZ -> return ShNil
  SS n -> ShCons <$> peekShape ptr off n
                 <*> (fromIntegral <$> peekByteOff @Int64 ptr (off + (fromSNat n) * 8))

compile' :: SList (Const String) env -> Ex env t -> CompM CExpr
compile' env = \case
  EVar _ t i -> do
    let Const var = slistIdx env i
    incrementVarAlways "var" Increment t var
    return $ CELit var

  ELet _ rhs body -> do
    var <- compileAssign "" env rhs
    rete <- compile' (Const var `SCons` env) body
    incrementVarAlways "let" Decrement (typeOf rhs) var
    return rete

  EPair _ a b -> do
    name <- emitStruct (STPair (typeOf a) (typeOf b))
    e1 <- compile' env a
    e2 <- compile' env b
    return $ CEStruct name [("a", e1), ("b", e2)]

  EFst _ e -> do
    let STPair _ t2 = typeOf e
    e' <- compile' env e
    case incrementVar "fst" Decrement t2 of
      Nothing -> return $ CEProj e' "a"
      Just f -> do var <- genName
                   emit $ SVarDecl True (repSTy (typeOf e)) var e'
                   f (var ++ ".b")
                   return $ CEProj (CELit var) "a"

  ESnd _ e -> do
    let STPair t1 _ = typeOf e
    e' <- compile' env e
    case incrementVar "snd" Decrement t1 of
      Nothing -> return $ CEProj e' "b"
      Just f -> do var <- genName
                   emit $ SVarDecl True (repSTy (typeOf e)) var e'
                   f (var ++ ".a")
                   return $ CEProj (CELit var) "b"

  ENil _ -> do
    name <- emitStruct STNil
    return $ CEStruct name []

  EInl _ t e -> do
    name <- emitStruct (STEither (typeOf e) t)
    e1 <- compile' env e
    return $ CEStruct name [("tag", CELit "0"), ("l", e1)]

  EInr _ t e -> do
    name <- emitStruct (STEither t (typeOf e))
    e2 <- compile' env e
    return $ CEStruct name [("tag", CELit "1"), ("r", e2)]

  ECase _ (EOp _ OIf e) a b -> do
    e1 <- compile' env e
    (e2, stmts2) <- scope $ compile' (Const undefined `SCons` env) a  -- don't access that nil, stupid you
    (e3, stmts3) <- scope $ compile' (Const undefined `SCons` env) b
    retvar <- genName
    emit $ SVarDeclUninit (repSTy (typeOf a)) retvar
    emit $ SIf e1
             (stmts2 <> pure (SAsg retvar e2))
             (stmts3 <> pure (SAsg retvar e3))
    return (CELit retvar)

  ECase _ e a b -> do
    let STEither t1 t2 = typeOf e
    e1 <- compile' env e
    var <- genName
    -- I know those are not variable names, but it's fine, probably
    (e2, stmts2) <- scope $ compile' (Const (var ++ ".l") `SCons` env) a
    (e3, stmts3) <- scope $ compile' (Const (var ++ ".r") `SCons` env) b
    ((), stmtsRel1) <- scope $ incrementVarAlways "case1" Decrement t1 (var ++ ".l")
    ((), stmtsRel2) <- scope $ incrementVarAlways "case2" Decrement t2 (var ++ ".r")
    retvar <- genName
    emit $ SVarDeclUninit (repSTy (typeOf a)) retvar
    emit $ SBlock (pure (SVarDecl True (repSTy (typeOf e)) var e1)
                <> pure (SIf (CEBinop (CEProj (CELit var) "tag") "==" (CELit "0"))
                           (stmts2
                            <> stmtsRel1
                            <> pure (SAsg retvar e2))
                           (stmts3
                            <> stmtsRel2
                            <> pure (SAsg retvar e3))))
    return (CELit retvar)

  ENothing _ t -> do
    name <- emitStruct (STMaybe t)
    return $ CEStruct name [("tag", CELit "0")]

  EJust _ e -> do
    name <- emitStruct (STMaybe (typeOf e))
    e1 <- compile' env e
    return $ CEStruct name [("tag", CELit "1"), ("j", e1)]

  EMaybe _ a b e -> do
    let STMaybe t = typeOf e
    e1 <- compile' env e
    var <- genName
    (e2, stmts2) <- scope $ compile' env a
    (e3, stmts3) <- scope $ compile' (Const (var ++ ".j") `SCons` env) b
    ((), stmtsRel) <- scope $ incrementVarAlways "maybe" Decrement t (var ++ ".j")
    retvar <- genName
    emit $ SVarDeclUninit (repSTy (typeOf a)) retvar
    emit $ SBlock (pure (SVarDecl True (repSTy (typeOf e)) var e1)
                <> pure (SIf (CEBinop (CEProj (CELit var) "tag") "==" (CELit "0"))
                           (stmts2
                            <> pure (SAsg retvar e2))
                           (stmts3
                            <> stmtsRel
                            <> pure (SAsg retvar e3))))
    return (CELit retvar)

  EConstArr _ n t (Array sh vec) -> do
    strname <- emitStruct (STArr n (STScal t))
    tldname <- genName' "carraybuf"
    -- Give it a refcount of _half_ the size_t max, so that it can be
    -- incremented and decremented at will and will "never" reach anything
    -- where something happens
    emitTLD $ "static " ++ strname ++ "_buf " ++ tldname ++ " = " ++
              "(" ++ strname ++ "_buf){.sh = {" ++ intercalate "," (map show (shapeToList sh)) ++ "}, " ++
              ".refc = (size_t)1<<63, .xs = {" ++ intercalate "," (map (compileScal False t) (toList vec)) ++ "}};"
    return (CEStruct strname [("buf", CEAddrOf (CELit tldname))])

  EBuild _ n esh efun -> do
    shname <- compileAssign "sh" env esh

    arrname <- allocArray "build" Malloc "arr" n (typeOf efun) Nothing (indexTupleComponents n shname)

    idxargname <- genName' "ix"
    (funretval, funstmts) <- scope $ compile' (Const idxargname `SCons` env) efun

    linivar <- genName' "li"
    ivars <- replicateM (fromSNat n) (genName' "i")
    emit $ SBlock $
      pure (SVarDecl False "size_t" linivar (CELit "0"))
      <> compose [pure . SLoop (repSTy tIx) ivar (CELit "0")
                               (CECast (repSTy tIx) (CEIndex (CELit (arrname ++ ".buf->sh")) (CELit (show dimidx))))
                 | (ivar, dimidx) <- zip ivars [0::Int ..]]
           (pure (SVarDecl True (repSTy (typeOf esh)) idxargname
                           (shapeTupFromLitVars n ivars))
            <> funstmts
            <> pure (SAsg (arrname ++ ".buf->xs[" ++ linivar ++ "++]") funretval))

    return (CELit arrname)

  EFold1Inner _ commut efun ex0 earr -> do
    let STArr (SS n) t = typeOf earr

    -- let vecwid = case commut of Commut -> 8 :: Int
    --                             Noncommut -> 1

    x0name <- compileAssign "foldx0" env ex0
    arrname <- compileAssign "foldarr" env earr

    zeroRefcountCheck (typeOf earr) "fold1i" arrname

    shszname <- genName' "shsz"
    -- This n is one less than the shape of the thing we're querying, which is
    -- unexpected. But it's exactly what we want, so we do it anyway.
    emit $ SVarDecl True (repSTy tIx) shszname (compileArrShapeSize n arrname)

    resname <- allocArray "fold" Malloc "foldres" n t (Just (CELit shszname))
                  [CELit (arrname ++ ".buf->sh[" ++ show i ++ "]") | i <- [0 .. fromSNat n - 1]]

    lenname <- genName' "n"
    emit $ SVarDecl True (repSTy tIx) lenname
                    (CELit (arrname ++ ".buf->sh[" ++ show (fromSNat n) ++ "]"))

    ((), x0incrStmts) <- scope $ incrementVarAlways "foldx0" Increment t x0name

    ivar <- genName' "i"
    jvar <- genName' "j"
    -- kvar <- if vecwid > 1 then genName' "k" else return ""

    accvar <- genName' "tot"
    let arreltlit = arrname ++ ".buf->xs[" ++ lenname ++ " * " ++ ivar ++ " + " ++
                    ({- if vecwid > 1 then show vecwid ++ " * " ++ jvar ++ " + " ++ kvar else -} jvar) ++ "]"
    (funres, funStmts) <- scope $ compile' (Const arreltlit `SCons` Const accvar `SCons` env) efun
    ((), arreltIncrStmts) <- scope $ incrementVarAlways "foldelt" Increment t arreltlit

    emit $ SLoop (repSTy tIx) ivar (CELit "0") (CELit shszname) $
             pure (SVarDecl False (repSTy t) accvar (CELit x0name))
             <> x0incrStmts  -- we're copying x0 here
             <> (pure $ SLoop (repSTy tIx) jvar (CELit "0") (CELit lenname) $
                   -- The combination function will consume the array element
                   -- and the accumulator. The accumulator is replaced by
                   -- what comes out of the function anyway, so that's
                   -- fine, but we do need to increment the array element.
                   arreltIncrStmts
                   <> funStmts
                   <> pure (SAsg accvar funres))
             <> pure (SAsg (resname ++ ".buf->xs[" ++ ivar ++ "]") (CELit accvar))

    incrementVarAlways "foldx0" Decrement t x0name
    incrementVarAlways "foldarr" Decrement (typeOf earr) arrname

    return (CELit resname)

  ESum1Inner _ e -> do
    let STArr (SS n) t = typeOf e
    argname <- compileAssign "sumarg" env e

    zeroRefcountCheck (typeOf e) "sum1i" argname

    shszname <- genName' "shsz"
    -- This n is one less than the shape of the thing we're querying, like EFold1Inner.
    emit $ SVarDecl True (repSTy tIx) shszname (compileArrShapeSize n argname)

    resname <- allocArray "sum" Malloc "sumres" n t (Just (CELit shszname))
                  [CELit (argname ++ ".buf->sh[" ++ show i ++ "]") | i <- [0 .. fromSNat n - 1]]

    lenname <- genName' "n"
    emit $ SVarDecl True (repSTy tIx) lenname
                    (CELit (argname ++ ".buf->sh[" ++ show (fromSNat n) ++ "]"))

    let vecwid = 8 :: Int
    ivar <- genName' "i"
    jvar <- genName' "j"
    kvar <- genName' "k"
    accvar <- genName' "tot"
    let nchunks = CEBinop (CELit lenname) "/" (CELit (show vecwid))
    emit $ SLoop (repSTy tIx) ivar (CELit "0") (CELit shszname) $ BList
             -- we have ScalIsNumeric, so it has 0 and (+) in C
             [SVerbatim $ repSTy t ++ " " ++ accvar ++ "[" ++ show vecwid ++ "] = {" ++ intercalate "," (replicate vecwid "0") ++ "};"
             ,SLoop (repSTy tIx) jvar (CELit "0") nchunks $
                pure $ SLoop (repSTy tIx) kvar (CELit "0") (CELit (show vecwid)) $
                  pure $ SVerbatim $ accvar ++ "[" ++ kvar ++ "] += " ++ argname ++ ".buf->xs[" ++ lenname ++ " * " ++ ivar ++ " + " ++ show vecwid ++ " * " ++ jvar ++ " + " ++ kvar ++ "];"
             ,SLoop (repSTy tIx) kvar (CELit "1") (CELit (show vecwid)) $
                pure $ SVerbatim $ accvar ++ "[0] += " ++ accvar ++ "[" ++ kvar ++ "];"
             ,SLoop (repSTy tIx) kvar (CEBinop nchunks "*" (CELit (show vecwid))) (CELit lenname) $
                pure $ SVerbatim $ accvar ++ "[0] += " ++ argname ++ ".buf->xs[" ++ lenname ++ " * " ++ ivar ++ " + " ++ kvar ++ "];"
             ,SAsg (resname ++ ".buf->xs[" ++ ivar ++ "]") (CELit (accvar++"[0]"))]

    incrementVarAlways "sum" Decrement (typeOf e) argname

    return (CELit resname)

  EUnit _ e -> do
    e' <- compile' env e
    let typ = STArr SZ (typeOf e)
    strname <- emitStruct typ
    name <- genName
    emit $ SVarDecl True strname name (CEStruct strname
                  [("buf", CECall "malloc_instr" [CELit (show (8 + sizeofSTy (typeOf e)))])])
    emit $ SAsg (name ++ ".buf->refc") (CELit "1")
    emit $ SAsg (name ++ ".buf->xs[0]") e'
    return (CELit name)

  EReplicate1Inner _ elen earg -> do
    let STArr n t = typeOf earg
    lenname <- compileAssign "replen" env elen
    argname <- compileAssign "reparg" env earg

    zeroRefcountCheck (typeOf earg) "replicate1i" argname

    shszname <- genName' "shsz"
    emit $ SVarDecl True (repSTy tIx) shszname (compileArrShapeSize n argname)

    resname <- allocArray "repl1i" Malloc "rep" (SS n) t
                 (Just (CEBinop (CELit shszname) "*" (CELit lenname)))
                 ([CELit (argname ++ ".buf->sh[" ++ show i ++ "]") | i <- [0 .. fromSNat n - 1]]
                  ++ [CELit lenname])

    ivar <- genName' "i"
    jvar <- genName' "j"
    emit $ SLoop (repSTy tIx) ivar (CELit "0") (CELit shszname) $
             pure $ SLoop (repSTy tIx) jvar (CELit "0") (CELit lenname) $
               pure $ SAsg (resname ++ ".buf->xs[" ++ ivar ++ " * " ++ lenname ++ " + " ++ jvar ++ "]")
                           (CELit (argname ++ ".buf->xs[" ++ ivar ++ "]"))

    incrementVarAlways "repl1i" Decrement (typeOf earg) argname

    return (CELit resname)

  EMaximum1Inner _ e -> compileExtremum "max" "maximum1i" ">" env e

  EMinimum1Inner _ e -> compileExtremum "min" "minimum1i" "<" env e

  EConst _ t x -> return $ CELit $ compileScal True t x

  EIdx0 _ e -> do
    let STArr _ t = typeOf e
    arrname <- compileAssign "" env e
    zeroRefcountCheck (typeOf e) "idx0" arrname
    name <- genName
    emit $ SVarDecl True (repSTy t) name
              (CEIndex (CEPtrProj (CEProj (CELit arrname) "buf") "xs") (CELit "0"))
    incrementVarAlways "idx0" Decrement (STArr SZ t) arrname
    return (CELit name)

  -- EIdx1 _ a b -> error "TODO" -- EIdx1 ext (compile' a) (compile' b)

  EIdx _ earr eidx -> do
    let STArr n t = typeOf earr
    arrname <- compileAssign "ixarr" env earr
    zeroRefcountCheck (typeOf earr) "idx" arrname
    idxname <- if fromSNat n > 0  -- prevent an unused-varable warning
                 then compileAssign "ixix" env eidx
                 else return ""  -- won't be used in this case

    when emitChecks $
      forM_ (zip [0::Int ..] (indexTupleComponents n idxname)) $ \(i, ixcomp) ->
        emit $ SIf (CEBinop (CEBinop ixcomp "<" (CELit "0")) "||"
                            (CEBinop ixcomp ">=" (CECast (repSTy tIx) (CELit (arrname ++ ".buf->sh[" ++ show i ++ "]")))))
                 (pure $ SVerbatim $
                    "fprintf(stderr, \"[chad-kernel] CHECK: index out of range (arr=%p)\\n\", " ++
                      arrname ++ ".buf); abort();")
                 mempty

    resname <- genName' "ixres"
    emit $ SVarDecl True (repSTy t) resname (CEIndex (CELit (arrname ++ ".buf->xs")) (toLinearIdx n arrname idxname))
    incrementVarAlways "idxelt" Increment t resname
    incrementVarAlways "idx" Decrement (STArr n t) arrname
    return (CELit resname)

  EShape _ e -> do
    let STArr n _ = typeOf e
        t = tTup (sreplicate n tIx)
    _ <- emitStruct t
    name <- compileAssign "" env e
    zeroRefcountCheck (typeOf e) "shape" name
    resname <- genName
    emit $ SVarDecl True (repSTy t) resname (compileShapeQuery n name)
    incrementVarAlways "shape" Decrement (typeOf e) name
    return (CELit resname)

  EOp _ op (EPair _ e1 e2) -> do
    e1' <- compile' env e1
    e2' <- compile' env e2
    compileOpPair op e1' e2'

  EOp _ op e -> do
    e' <- compile' env e
    compileOpGeneral op e'

  ECustom _ _ _ _ earg _ _ e1 e2 -> do
    name1 <- compileAssign "" env e1
    name2 <- compileAssign "" env e2
    case (incrementVar "custom1" Decrement (typeOf e1), incrementVar "custom2" Decrement (typeOf e2)) of
      (Nothing, Nothing) -> compile' (Const name2 `SCons` Const name1 `SCons` SNil) earg
      (mfun1, mfun2) -> do
        name <- compileAssign "" (Const name2 `SCons` Const name1 `SCons` SNil) earg
        maybe (return ()) ($ name1) mfun1
        maybe (return ()) ($ name2) mfun2
        return (CELit name)

  EWith _ t e1 e2 -> do
    actyname <- emitStruct (STAccum t)
    name1 <- compileAssign "" env e1

    zeroRefcountCheck (typeOf e1) "with" name1

    emit $ SVerbatim $ "// copyForWriting start (" ++ name1 ++ ")"
    mcopy <- copyForWriting (CHAD.d2 t) name1
    accname <- genName' "accum"
    emit $ SVarDecl False actyname accname
              (CEStruct actyname [("buf", CECall "malloc_instr" [CELit (show (sizeofSTy (CHAD.d2 t)))])])
    emit $ SAsg (accname++".buf->ac") (maybe (CELit name1) id mcopy)
    emit $ SVerbatim $ "// initial accumulator constructed (" ++ name1 ++ ")."

    e2' <- compile' (Const accname `SCons` env) e2

    resname <- genName' "acret"
    emit $ SVarDecl True (repSTy (CHAD.d2 t)) resname (CELit (accname++".buf->ac"))
    emit $ SVerbatim $ "free_instr(" ++ accname ++ ".buf);"

    rettyname <- emitStruct (STPair (typeOf e2) (CHAD.d2 t))
    return $ CEStruct rettyname [("a", e2'), ("b", CELit resname)]

  EAccum _ t prj eidx eval eacc -> do
    nameidx <- compileAssign "acidx" env eidx
    nameval <- compileAssign "acval" env eval

    -- Generate the variable manually because this one has to be non-const.
    eacc' <- compile' env eacc
    nameacc <- genName' "acac"
    emit $ SVarDecl False (repSTy (typeOf eacc)) nameacc eacc'

    let -- Expects a variable reference to a value of type @D2 a@.
        setZero :: STy a -> String -> CompM ()
        setZero STNil _ = return ()
        setZero STPair{} v = emit $ SAsg (v++".tag") (CELit "0")  -- Maybe (Pair (D2 a) (D2 b))
        setZero STEither{} v = emit $ SAsg (v++".tag") (CELit "0")  -- Maybe (Either (D2 a) (D2 b))
        setZero STMaybe{} v = emit $ SAsg (v++".tag") (CELit "0")  -- Maybe (D2 a)
        setZero STArr{} v = emit $ SAsg (v++".tag") (CELit "0")  -- Maybe (Arr n (D2 a))
        setZero (STScal sty) v = case sty of
          STI32 -> return ()  -- Nil
          STI64 -> return ()  -- Nil
          STF32 -> emit $ SAsg v (CELit "0.0f")
          STF64 -> emit $ SAsg v (CELit "0.0")
          STBool -> return ()  -- Nil
        setZero STAccum{} _ = error "Compile: setZero: nested accumulators unsupported"

        initD2Pair :: STy a -> STy b -> String -> CompM ()
        initD2Pair a b v = do  -- Maybe (Pair (D2 a) (D2 b))
          ((), stmts1) <- scope $ setZero a (v++".j.a")
          ((), stmts2) <- scope $ setZero b (v++".j.b")
          emit $ SIf (CEBinop (CELit (v++".tag")) "==" (CELit "0"))
                   (pure (SAsg (v++".tag") (CELit "1")) <> stmts1 <> stmts2)
                   mempty

        initD2Either :: STy a -> STy b -> String -> Either () () -> CompM ()
        initD2Either a b v side = do  -- Maybe (Either (D2 a) (D2 b))
          ((), stmts) <- case side of
                           Left () -> scope $ setZero a (v++".j.l")
                           Right () -> scope $ setZero b (v++".j.r")
          emit $ SIf (CEBinop (CELit (v++".tag")) "==" (CELit "0"))
                   (pure (SAsg (v++".tag") (CELit "1")) <> stmts)
                   mempty

        initD2Maybe :: STy a -> String -> CompM ()
        initD2Maybe a v = do  -- Maybe (D2 a)
          ((), stmts) <- scope $ setZero a (v++".j")
          emit $ SIf (CEBinop (CELit (v++".tag")) "==" (CELit "0"))
                   (pure (SAsg (v++".tag") (CELit "1")) <> stmts)
                   mempty

    -- mind: this has to traverse the D2 of these things, and it also has to
    -- initialise data structures that are still sparse in the accumulator.
    let accumRef :: STy a -> SAcPrj p a b -> String -> String -> CompM String
        accumRef _ SAPHere v _ = pure v
        accumRef (STPair ta tb) (SAPFst prj') v i = do
          initD2Pair ta tb v
          accumRef ta prj' (v++".j.a") i
        accumRef (STPair ta tb) (SAPSnd prj') v i = do
          initD2Pair ta tb v
          accumRef tb prj' (v++".j.b") i
        accumRef (STEither ta tb) (SAPLeft prj') v i = do
          initD2Either ta tb v (Left ())
          accumRef ta prj' (v++".j.l") i
        accumRef (STEither ta tb) (SAPRight prj') v i = do
          initD2Either ta tb v (Right ())
          accumRef tb prj' (v++".j.r") i
        accumRef (STMaybe tj) (SAPJust prj') v i = do
          initD2Maybe tj v
          accumRef tj prj' (v++".j") i
        accumRef (STArr n t') (SAPArrIdx prj' _) v i = do
          (newarrName, newarrStmts) <- scope $ allocArray "accumRef" Calloc "prjarr" n t' Nothing (indexTupleComponents n (i++".a.b"))
          emit $ SIf (CEBinop (CELit (v++".tag")) "==" (CELit "0"))
                   (pure (SAsg (v++".tag") (CELit "1"))
                    <> newarrStmts
                    <> pure (SAsg (v++".j") (CELit newarrName)))
                   mempty

          when emitChecks $ do
            let shfmt = "[" ++ intercalate "," (replicate (fromSNat n) "%\"PRIi64\"") ++ "]"
            forM_ (zip3 [0::Int ..]
                        (indexTupleComponents n (i++".a.a"))
                        (indexTupleComponents n (i++".a.b"))) $ \(j, ixcomp, shcomp) -> do
              let a .||. b = CEBinop a "||" b
              emit $ SIf (CEBinop ixcomp "<" (CELit "0")
                          .||.
                          CEBinop ixcomp ">=" (CECast (repSTy tIx) (CELit (v ++ ".j.buf->sh[" ++ show j ++ "]")))
                          .||.
                          CEBinop shcomp "!=" (CECast (repSTy tIx) (CELit (v ++ ".j.buf->sh[" ++ show j ++ "]"))))
                       (pure $ SVerbatim $
                          "fprintf(stderr, \"[chad-kernel] CHECK: accum prj incorrect (arr=%p, " ++
                          "arrsh=" ++ shfmt ++ ", acix=" ++ shfmt ++ ", acsh=" ++ shfmt ++ ")\\n\", " ++
                          v ++ ".j.buf" ++
                          concat [", " ++ v ++ ".j.buf->sh[" ++ show k ++ "]" | k <- [0 .. fromSNat n - 1]] ++
                          concat [", " ++ printCExpr 2 comp "" | comp <- indexTupleComponents n (i++".a.a")] ++
                          concat [", " ++ printCExpr 2 comp "" | comp <- indexTupleComponents n (i++".a.b")] ++
                          "); " ++
                          "abort();")
                       mempty

          accumRef t' prj' (v++".j.buf->xs[" ++ printCExpr 0 (toLinearIdx n (v++".j") (i++".a.a")) "]") (i++".b")

    -- mind: this has to add the D2 of these things, and it also has to
    -- initialise data structures that are still sparse in the accumulator.
    let add :: STy a -> String -> String -> CompM ()
        add STNil _ _ = return ()
        add (STPair t1 t2) d s = do
          ((), stmts1) <- scope $ add t1 (d++".j.a") (s++".j.a")
          ((), stmts2) <- scope $ add t2 (d++".j.b") (s++".j.b")
          emit $ SIf (CEBinop (CELit (d++".tag")) "==" (CELit "0"))
                   (pure (SAsg d (CELit s)))
                   (pure (SIf (CEBinop (CELit (s++".tag")) "==" (CELit "1"))
                            (stmts1 <> stmts2)
                            mempty))
        add (STEither t1 t2) d s = do
          ((), stmts1) <- scope $ add t1 (d++".j.l") (s++".j.l")
          ((), stmts2) <- scope $ add t2 (d++".j.r") (s++".j.r")
          emit $ SIf (CEBinop (CELit (d++".tag")) "==" (CELit "0"))
                   (pure (SAsg d (CELit s)))
                   (pure (SIf (CEBinop (CELit (s++".tag")) "==" (CELit "1"))
                            (pure (SAsg (d++".j.tag") (CELit (s++".j.tag")))
                             <> pure (SIf (CEBinop (CELit (s++".j.tag")) "==" (CELit "0"))
                                        stmts1 stmts2))
                            mempty))
        add (STMaybe t1) d s = do
          ((), stmts1) <- scope $ add t1 (d++".j") (s++".j")
          emit $ SIf (CEBinop (CELit (d++".tag")) "==" (CELit "0"))
                   (pure (SAsg d (CELit s)))
                   (pure (SIf (CEBinop (CELit (s++".tag")) "==" (CELit "1"))
                            (pure (SAsg (d++".tag") (CELit "1")) <> stmts1)
                            mempty))
        add (STArr n t1) d s = do
          shsizename <- genName' "acshsz"
          ivar <- genName' "i"
          ((), stmts1) <- scope $ add t1 (d++".j.buf->xs["++ivar++"]") (s++".j.buf->xs["++ivar++"]")
          ((), stmtsDecr) <- scope $ incrementVarAlways "accumarr" Decrement (STArr n (CHAD.d2 t1)) (s++".j")
          emit $ SIf (CEBinop (CELit (s++".tag")) "==" (CELit "1"))
                   (pure (SIf (CEBinop (CELit (d++".tag")) "==" (CELit "0"))
                            (pure (SAsg d (CELit s)))
                            (pure (SVarDecl True (repSTy tIx) shsizename (compileArrShapeSize n (s++".j")))
                             -- TODO: emit check here for the source being either equal in shape to the destination
                             <> pure (SLoop (repSTy tIx) ivar (CELit "0") (CELit shsizename)
                                        stmts1)
                             <> stmtsDecr)))
                   mempty
        add (STScal sty) d s = case sty of
          STI32 -> return ()
          STI64 -> return ()
          STF32 -> emit $ SVerbatim $ d ++ " += " ++ s ++ ";"
          STF64 -> emit $ SVerbatim $ d ++ " += " ++ s ++ ";"
          STBool -> return ()
        add (STAccum _) _ _ = error "Compile: nested accumulators unsupported"

    emit $ SVerbatim $ "// compile EAccum start (" ++ show prj ++ ")"
    dest <- accumRef t prj (nameacc++".buf->ac") nameidx
    add (acPrjTy prj t) dest nameval
    emit $ SVerbatim $ "// compile EAccum end"

    return $ CEStruct (repSTy STNil) []

  EError _ t s -> do
    let padleft len c s' = replicate (len - length s) c ++ s'
        escape = concatMap $ \c -> if | c `elem` "\"\\" -> ['\\',c]
                                      | ord c < 32      -> "\\x" ++ padleft 2 '0' (showHex (ord c) "")
                                      | otherwise       -> [c]
    emit $ SVerbatim $ "fputs(\"ERROR: " ++ escape s ++ "\\n\", stderr); exit(1);"
    case t of
      STScal _ -> return (CELit "0")
      _ -> do
        name <- emitStruct t
        return $ CEStruct name []

  EZero{} -> error "Compile: monoid operations should have been eliminated"
  EPlus{} -> error "Compile: monoid operations should have been eliminated"
  EOneHot{} -> error "Compile: monoid operations should have been eliminated"

  EIdx1{} -> error "Compile: not implemented: EIdx1"

compileAssign :: String -> SList (Const String) env -> Ex env t -> CompM String
compileAssign prefix env e = do
  e' <- compile' env e
  name <- genName' prefix
  emit $ SVarDecl True (repSTy (typeOf e)) name e'
  return name

data Increment = Increment | Decrement
  deriving (Show)

-- | Increment reference counts in the components of the given variable.
incrementVar :: String -> Increment -> STy a -> Maybe (String -> CompM ())
incrementVar marker inc ty =
  let tree = makeArrayTree ty
  in case tree of ATNoop -> Nothing
                  _ -> Just $ \var -> incrementVar' marker inc var tree

incrementVarAlways :: String -> Increment -> STy a -> String -> CompM ()
incrementVarAlways marker inc ty var = maybe (pure ()) ($ var) (incrementVar marker inc ty)

data ArrayTree = ATArray (Some SNat) (Some STy)  -- ^ we've arrived at an array we need to decrement the refcount of (contains rank and element type of the array)
               | ATNoop  -- ^ don't do anything here
               | ATProj String ArrayTree  -- ^ descend one field deeper
               | ATCondTag ArrayTree ArrayTree  -- ^ if tag is 0, first; if 1, second
               | ATBoth ArrayTree ArrayTree  -- ^ do both these paths

smartATProj :: String -> ArrayTree -> ArrayTree
smartATProj _ ATNoop = ATNoop
smartATProj field t = ATProj field t

smartATCondTag :: ArrayTree -> ArrayTree -> ArrayTree
smartATCondTag ATNoop ATNoop = ATNoop
smartATCondTag t t' = ATCondTag t t'

smartATBoth :: ArrayTree -> ArrayTree -> ArrayTree
smartATBoth ATNoop t = t
smartATBoth t ATNoop = t
smartATBoth t t' = ATBoth t t'

makeArrayTree :: STy a -> ArrayTree
makeArrayTree STNil = ATNoop
makeArrayTree (STPair a b) = smartATBoth (smartATProj "a" (makeArrayTree a))
                                         (smartATProj "b" (makeArrayTree b))
makeArrayTree (STEither a b) = smartATCondTag (smartATProj "l" (makeArrayTree a))
                                              (smartATProj "r" (makeArrayTree b))
makeArrayTree (STMaybe t) = smartATCondTag ATNoop (smartATProj "j" (makeArrayTree t))
makeArrayTree (STArr n t) = ATArray (Some n) (Some t)
makeArrayTree (STScal _) = ATNoop
makeArrayTree (STAccum _) = ATNoop

incrementVar' :: String -> Increment -> String -> ArrayTree -> CompM ()
incrementVar' marker inc path (ATArray (Some n) (Some eltty)) =
  case inc of
    Increment -> do
      emit $ SVerbatim (path ++ ".buf->refc++;")
      when debugRefc $
        emit $ SVerbatim $ "fprintf(stderr, \"[chad-kernel] arr %p in+ -> %zu <" ++ marker ++ ">\\n\", " ++ path ++ ".buf, " ++ path ++ ".buf->refc);"
    Decrement -> do
      case incrementVar (marker++".elt") Decrement eltty of
        Nothing ->
          if debugRefc
            then do
              emit $ SVerbatim $ "fprintf(stderr, \"[chad-kernel] arr %p de- -> %zu <" ++ marker ++ ">\", " ++ path ++ ".buf, " ++ path ++ ".buf->refc - 1);"
              emit $ SVerbatim $ "if (--" ++ path ++ ".buf->refc == 0) { fprintf(stderr, \"; free(\"); free_instr(" ++ path ++ ".buf); fprintf(stderr, \") ok\\n\"); } else fprintf(stderr, \"\\n\");"
            else do
              emit $ SVerbatim $ "if (--" ++ path ++ ".buf->refc == 0) free_instr(" ++ path ++ ".buf);"
        Just f -> do
          when debugRefc $
            emit $ SVerbatim $ "fprintf(stderr, \"[chad-kernel] arr %p de- -> %zu <" ++ marker ++ "> recfree\\n\", " ++ path ++ ".buf, " ++ path ++ ".buf->refc - 1);"
          shszvar <- genName' "frshsz"
          ivar <- genName' "i"
          ((), eltDecrStmts) <- scope $ f (path ++ ".buf->xs[" ++ ivar ++ "]")
          emit $ SIf (CELit ("--" ++ path ++ ".buf->refc == 0"))
                   (BList [SVarDecl True "size_t" shszvar (compileArrShapeSize n path)
                          ,SLoop "size_t" ivar (CELit "0") (CELit shszvar) $
                             eltDecrStmts
                          ,SVerbatim $ "free_instr(" ++ path ++ ".buf);"])
                   mempty
incrementVar' _ _ _ ATNoop = pure ()
incrementVar' marker inc path (ATProj field t) = incrementVar' (marker++"."++field) inc (path ++ "." ++ field) t
incrementVar' marker inc path (ATCondTag t1 t2) = do
  ((), stmts1) <- scope $ incrementVar' (marker++".t1") inc path t1
  ((), stmts2) <- scope $ incrementVar' (marker++".t2") inc path t2
  emit $ SIf (CEBinop (CELit (path ++ ".tag")) "==" (CELit "0")) stmts1 stmts2
incrementVar' marker inc path (ATBoth t1 t2) = incrementVar' (marker++".1") inc path t1 >> incrementVar' (marker++".2") inc path t2

toLinearIdx :: SNat n -> String -> String -> CExpr
toLinearIdx SZ _ _ = CELit "0"
toLinearIdx (SS SZ) _ idxvar = CELit (idxvar ++ ".b")
toLinearIdx (SS n) arrvar idxvar =
  CEBinop (CEBinop (toLinearIdx n arrvar (idxvar ++ ".a"))
                   "*" (CEIndex (CELit (arrvar ++ ".buf->sh")) (CELit (show (fromSNat n)))))
          "+" (CELit (idxvar ++ ".b"))

-- fromLinearIdx :: SNat n -> String -> String -> CompM CExpr
-- fromLinearIdx SZ _ _ = return $ CEStruct (repSTy STNil) []
-- fromLinearIdx (SS n) arrvar idxvar = do
--   name <- genName
--   emit $ SVarDecl True (repSTy tIx) name (CEBinop (CELit idxvar) "/" (CELit (arrvar ++ ".buf->sh[" ++ show (fromSNat n) ++ "]")))
--   _

data AllocMethod = Malloc | Calloc
  deriving (Show)

-- | The shape must have the outer dimension at the head (and the inner dimension on the right).
allocArray :: String -> AllocMethod -> String -> SNat n -> STy t -> Maybe CExpr -> [CExpr] -> CompM String
allocArray marker method nameBase rank eltty mshsz shape = do
  when (length shape /= fromSNat rank) $
    error "allocArray: shape does not match rank"
  let arrty = STArr rank eltty
  strname <- emitStruct arrty
  arrname <- genName' nameBase
  shsz <- case mshsz of
            Just e -> return e
            Nothing -> return (foldl0' (\a b -> CEBinop a "*" b) (CELit "1") shape)
  let nbytesExpr = CEBinop (CELit (show (fromSNat rank * 8 + 8)))
                           "+"
                           (CEBinop shsz "*" (CELit (show (sizeofSTy eltty))))
  emit $ SVarDecl True strname arrname $ CEStruct strname
            [("buf", case method of Malloc -> CECall "malloc_instr" [nbytesExpr]
                                    Calloc -> CECall "calloc_instr" [nbytesExpr])]
  forM_ (zip shape [0::Int ..]) $ \(dim, i) ->
    emit $ SAsg (arrname ++ ".buf->sh[" ++ show i ++ "]") dim
  emit $ SAsg (arrname ++ ".buf->refc") (CELit "1")
  when debugRefc $
    emit $ SVerbatim $ "fprintf(stderr, \"[chad-kernel] arr %p allocated <" ++ marker ++ ">\\n\", " ++ arrname ++ ".buf);"
  return arrname

compileShapeQuery :: SNat n -> String -> CExpr
compileShapeQuery SZ _ = CEStruct (repSTy STNil) []
compileShapeQuery (SS n) var =
  CEStruct (repSTy (tTup (sreplicate (SS n) tIx)))
    [("a", compileShapeQuery n var)
    ,("b", CEIndex (CELit (var ++ ".buf->sh")) (CELit (show (fromSNat n))))]

-- | Takes a variable name for the array, not the buffer.
compileArrShapeSize :: SNat n -> String -> CExpr
compileArrShapeSize SZ _ = CELit "1"
compileArrShapeSize n var =
  foldl1' (\a b -> CEBinop a "*" b) [CELit (var ++ ".buf->sh[" ++ show i ++ "]")
                                    | i <- [0 .. fromSNat n - 1]]

indexTupleComponents :: SNat n -> String -> [CExpr]
indexTupleComponents = \n var -> map CELit (toList (go n var))
  where
    go :: SNat n -> String -> Bag String
    go SZ _ = mempty
    go (SS n) var = go n (var ++ ".a") <> pure (var ++ ".b")

-- | Takes variable names with the innermost dimension on the right.
shapeTupFromLitVars :: SNat n -> [String] -> CExpr
shapeTupFromLitVars = \n -> go n . reverse
  where
    -- takes variables with the innermost dimension at the _head_
    go :: SNat n -> [String] -> CExpr
    go SZ [] = CEStruct (repSTy STNil) []
    go (SS n) (var : vars) = CEStruct (repSTy (tTup (sreplicate (SS n) tIx))) [("a", go n vars), ("b", CELit var)]
    go _ _ = error "shapeTupFromLitVars: SNat and list do not correspond"

compileOpGeneral :: SOp a b -> CExpr -> CompM CExpr
compileOpGeneral op e1 = do
  let unary cop = return @CompM $ CECall cop [e1]
  let binary cop = do
        name <- genName
        emit $ SVarDecl True (repSTy (opt1 op)) name e1
        return $ CEBinop (CEProj (CELit name) "a") cop (CEProj (CELit name) "b")
  case op of
    OAdd _ -> binary "+"
    OMul _ -> binary "*"
    ONeg _ -> unary "-"
    OLt _ -> binary "<"
    OLe _ -> binary "<="
    OEq _ -> binary "=="
    ONot -> unary "!"
    OAnd -> binary "&&"
    OOr -> binary "||"
    OIf -> do
      name <- emitStruct (STEither STNil STNil)
      _ <- emitStruct STNil
      return $ CEIf e1 (CEStruct name [("tag", CELit "0")])
                       (CEStruct name [("tag", CELit "1")])
    ORound64 -> unary "(int64_t)round"  -- ew
    OToFl64 -> unary "(double)"
    ORecip _ -> return $ CEBinop (CELit "1.0") "/" e1
    OExp STF32 -> unary "expf"
    OExp STF64 -> unary "exp"
    OLog STF32 -> unary "logf"
    OLog STF64 -> unary "log"
    OIDiv _ -> binary "/"

compileOpPair :: SOp a b -> CExpr -> CExpr -> CompM CExpr
compileOpPair op e1 e2 = do
  let binary cop = return @CompM $ CEBinop e1 cop e2
  case op of
    OAdd _ -> binary "+"
    OMul _ -> binary "*"
    OLt _ -> binary "<"
    OLe _ -> binary "<="
    OEq _ -> binary "=="
    OAnd -> binary "&&"
    OOr -> binary "||"
    OIDiv _ -> binary "/"
    _ -> error "compileOpPair: got unary operator"

-- | Bool: whether to ensure that the literal itself already has the appropriate type
compileScal :: Bool -> SScalTy t -> ScalRep t -> String
compileScal pedantic typ x = case typ of
  STI32 -> (if pedantic then "(int32_t)" else "") ++ show x
  STI64 -> (if pedantic then "(int64_t)" else "") ++ show x
  STF32 -> show x ++ "f"
  STF64 -> show x
  STBool -> if x then "1" else "0"

compileExtremum :: String -> String -> String -> SList (Const String) env -> Ex env (TArr (S n) t) -> CompM CExpr
compileExtremum nameBase opName operator env e = do
  let STArr (SS n) t = typeOf e
  argname <- compileAssign (nameBase ++ "arg") env e

  zeroRefcountCheck (typeOf e) opName argname

  shszname <- genName' "shsz"
  -- This n is one less than the shape of the thing we're querying, which is
  -- unexpected. But it's exactly what we want, so we do it anyway.
  emit $ SVarDecl True (repSTy tIx) shszname (compileArrShapeSize n argname)

  resname <- allocArray nameBase Malloc (nameBase ++ "res") n t (Just (CELit shszname))
                [CELit (argname ++ ".buf->sh[" ++ show i ++ "]") | i <- [0 .. fromSNat n - 1]]

  lenname <- genName' "n"
  emit $ SVarDecl True (repSTy tIx) lenname
                  (CELit (argname ++ ".buf->sh[" ++ show (fromSNat n) ++ "]"))

  emit $ SVerbatim $ "if (" ++ lenname ++ " == 0) { fprintf(stderr, \"Empty array in " ++ opName ++ "\\n\"); abort(); }"

  ivar <- genName' "i"
  jvar <- genName' "j"
  xvar <- genName
  redvar <- genName' "red"  -- use "red", not "acc", to avoid confusion with accumulators
  emit $ SLoop (repSTy tIx) ivar (CELit "0") (CELit shszname) $ BList
           -- we have ScalIsNumeric, so it has 1 and (<) etc. in C
           [SVarDecl False (repSTy t) redvar (CELit (argname ++ ".buf->xs[" ++ lenname ++ " * " ++ ivar ++ "]"))
           ,SLoop (repSTy tIx) jvar (CELit "1") (CELit lenname) $ BList
              [SVarDecl True (repSTy t) xvar (CELit (argname ++ ".buf->xs[" ++ lenname ++ " * " ++ ivar ++ " + " ++ jvar ++ "]"))
              ,SAsg redvar $ CEIf (CEBinop (CELit xvar) operator (CELit redvar)) (CELit xvar) (CELit redvar)
              ]
           ,SAsg (resname ++ ".buf->xs[" ++ ivar ++ "]") (CELit redvar)]

  incrementVarAlways nameBase Decrement (typeOf e) argname

  return (CELit resname)

-- | If this returns Nothing, there was nothing to copy because making a simple
-- value copy in C already makes it suitable to write to.
copyForWriting :: STy t -> String -> CompM (Maybe CExpr)
copyForWriting topty var = case topty of
  STNil -> return Nothing

  STPair a b -> do
    e1 <- copyForWriting a (var ++ ".a")
    e2 <- copyForWriting b (var ++ ".b")
    case (e1, e2) of
      (Nothing, Nothing) -> return Nothing
      _ -> return $ Just $ CEStruct (repSTy topty)
                             [("a", fromMaybe (CELit (var++".a")) e1)
                             ,("b", fromMaybe (CELit (var++".b")) e2)]

  STEither a b -> do
    (e1, stmts1) <- scope $ copyForWriting a (var ++ ".l")
    (e2, stmts2) <- scope $ copyForWriting b (var ++ ".r")
    case (e1, e2) of
      (Nothing, Nothing) -> return Nothing
      _ -> do
        name <- genName
        emit $ SVarDeclUninit (repSTy topty) name
        emit $ SIf (CEBinop (CELit (var++".tag")) "==" (CELit "0"))
                 (stmts1
                  <> pure (SAsg name (CEStruct (repSTy topty)
                                        [("tag", CELit "0"), ("l", fromMaybe (CELit (var++".l")) e1)])))
                 (stmts2
                  <> pure (SAsg name (CEStruct (repSTy topty)
                                        [("tag", CELit "1"), ("r", fromMaybe (CELit (var++".r")) e2)])))
        return (Just (CELit name))

  STMaybe t -> do
    (e1, stmts1) <- scope $ copyForWriting t (var ++ ".j")
    case e1 of
      Nothing -> return Nothing
      Just e1' -> do
        name <- genName
        emit $ SVarDeclUninit (repSTy topty) name
        emit $ SIf (CEBinop (CELit (var++".tag")) "==" (CELit "0"))
                 (pure (SAsg name (CEStruct (repSTy topty) [("tag", CELit "0")])))
                 (stmts1
                  <> pure (SAsg name (CEStruct (repSTy topty) [("tag", CELit "1"), ("j", e1')])))
        return (Just (CELit name))

  -- If there are no nested arrays, we know that a refcount of 1 means that the
  -- whole thing is owned. Nested arrays have their own refcount, so with
  -- nesting we'd have to check the refcounts of all the nested arrays _too_;
  -- let's not do that. Furthermore, no sub-arrays means that the whole thing
  -- is flat, and we can just memcpy if necessary.
  STArr n t | not (hasArrays t) -> do
    name <- genName
    shszname <- genName' "shsz"
    emit $ SVarDeclUninit (repSTy (STArr n t)) name

    when debugShapes $ do
      let shfmt = "[" ++ intercalate "," (replicate (fromSNat n) "%\"PRIi64\"") ++ "]"
      emit $ SVerbatim $
        "fprintf(stderr, \"[chad-kernel] with array " ++ shfmt ++ "\\n\"" ++
        concat [", " ++ var ++ ".buf->sh[" ++ show i ++ "]" | i <- [0 .. fromSNat n - 1]] ++
        ");"

    emit $ SIf (CEBinop (CELit (var ++ ".buf->refc")) "==" (CELit "1"))
             (pure (SAsg name (CELit var)))
             (let shbytes = fromSNat n * 8
                  databytes = CEBinop (CELit shszname) "*" (CELit (show (sizeofSTy t)))
                  totalbytes = CEBinop (CELit (show (shbytes + 8))) "+" databytes
              in BList
               [SVarDecl True (repSTy tIx) shszname (compileArrShapeSize n var)
               ,SAsg name (CEStruct (repSTy (STArr n t)) [("buf", CECall "malloc_instr" [totalbytes])])
               ,SVerbatim $ "memcpy(" ++ name ++ ".buf->sh, " ++ var ++ ".buf->sh, " ++
                                      show shbytes ++ ");"
               ,SAsg (name ++ ".buf->refc") (CELit "1")
               ,SVerbatim $ "memcpy(" ++ name ++ ".buf->xs, " ++ var ++ ".buf->xs, " ++
                                      printCExpr 0 databytes ");"])
    return (Just (CELit name))

  STArr n t -> do
    shszname <- genName' "shsz"
    emit $ SVarDecl True (repSTy tIx) shszname (compileArrShapeSize n var)

    let shbytes = fromSNat n * 8
        databytes = CEBinop (CELit shszname) "*" (CELit (show (sizeofSTy t)))
        totalbytes = CEBinop (CELit (show (shbytes + 8))) "+" databytes

    name <- genName
    emit $ SVarDecl False (repSTy (STArr n t)) name
             (CEStruct (repSTy (STArr n t)) [("buf", CECall "malloc_instr" [totalbytes])])
    emit $ SVerbatim $ "memcpy(" ++ name ++ ".buf->sh, " ++ var ++ ".buf->sh, " ++
                                 show shbytes ++ ");"
    emit $ SAsg (name ++ ".buf->refc") (CELit "1")

    -- put the arrays in variables to cut short the not-quite-var chain
    dstvar <- genName' "cpydst"
    emit $ SVarDecl True (repSTy t ++ " *") dstvar (CELit (name ++ ".buf->xs"))
    srcvar <- genName' "cpysrc"
    emit $ SVarDecl True (repSTy t ++ " *") srcvar (CELit (var ++ ".buf->xs"))

    ivar <- genName' "i"

    (cpye, cpystmts) <- scope $ copyForWriting t (srcvar ++ "[" ++ ivar ++ "]")
    let cpye' = case cpye of
                  Just e -> e
                  Nothing -> error "copyForWriting: arrays cannot be copied as-is, bug"

    emit $ SLoop (repSTy tIx) ivar (CELit "0") (CELit shszname) $
             cpystmts
             <> pure (SAsg (dstvar ++ "[" ++ ivar ++ "]") cpye')

    return (Just (CELit name))

  STScal _ -> return Nothing

  STAccum _ -> error "Compile: Nested accumulators not supported"

zeroRefcountCheck :: STy t -> String -> String -> CompM ()
zeroRefcountCheck toptyp opname topvar =
  when emitChecks $ do
    mstmts <- onlyIdGen $ runMaybeT (go toptyp topvar)
    case mstmts of
      Nothing -> return ()
      Just stmts -> forM_ stmts emit
  where
    -- | If this returns 'Nothing', no statements need to be generated for this type.
    go :: STy t -> String -> MaybeT IdGen.IdGen (Bag Stmt)
    go STNil _ = empty
    go (STPair a b) path = do
      (s1, s2) <- combine (go a (path++".a")) (go b (path++".b"))
      return (s1 <> s2)
    go (STEither a b) path = do
      (s1, s2) <- combine (go a (path++".l")) (go b (path++".r"))
      return $ pure $ SIf (CEBinop (CELit (path++".tag")) "==" (CELit "0")) s1 s2
    go (STMaybe a) path = do
      ss <- go a (path++".j")
      return $ pure $ SIf (CEBinop (CELit (path++".tag")) "==" (CELit "1")) ss mempty
    go (STArr n a) path = do
      ivar <- genName' "i"
      ss <- go a (path++".buf->xs["++ivar++"]")
      shszname <- genName' "shsz"
      let s1 = SVerbatim $
                 "if (__builtin_expect(" ++ path ++ ".buf->refc == 0, 0)) { " ++
                 "fprintf(stderr, \"[chad-kernel] CHECK: '" ++ opname ++ "' got array " ++
                 "%p with refc=0\\n\", " ++ path ++ ".buf); abort(); }"
      let s2 = SVarDecl True (repSTy tIx) shszname (compileArrShapeSize n path)
      let s3 = SLoop (repSTy tIx) ivar (CELit "0") (CELit shszname) ss
      return (BList [s1, s2, s3])
    go STScal{} _ = empty
    go STAccum{} _ = error "zeroRefcountCheck: passed an accumulator"

    combine :: (Monoid a, Monoid b, Monad m) => MaybeT m a -> MaybeT m b -> MaybeT m (a, b)
    combine (MaybeT a) (MaybeT b) = MaybeT $ do
      x <- a
      y <- b
      return $ case (x, y) of
        (Nothing, Nothing) -> Nothing
        (Just x', Nothing) -> Just (x', mempty)
        (Nothing, Just y') -> Just (mempty, y')
        (Just x', Just y') -> Just (x', y')

compose :: Foldable t => t (a -> a) -> a -> a
compose = foldr (.) id

showPtr :: Ptr a -> String
showPtr (Ptr a) = "0x" ++ showHex (integerFromWord# (int2Word# (addr2Int# a))) ""

-- | Type-restricted.
(^) :: Num a => a -> Int -> a
(^) = (Prelude.^)

foldl0' :: (a -> a -> a) -> a -> [a] -> a
foldl0' _ x [] = x
foldl0' f _ l = foldl1' f l