{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeApplications #-}
module Compile (compile) where

import Control.Monad (when)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.State.Strict
import Control.Monad.Trans.Writer.CPS
import Data.Bifunctor (first)
import Data.Foldable (toList)
import Data.Functor.Const
import qualified Data.Functor.Product as Product
import Data.Functor.Product (Product)
import Data.List (intersperse, intercalate)
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import Data.Set (Set)
import Data.Some
import qualified Data.Vector as V
import Foreign
import System.IO (hPutStrLn, stderr)

import Prelude hiding ((^))
import qualified Prelude

import Array
import AST
import AST.Pretty (ppTy)
import Compile.Exec
import Data
import Interpreter.Rep


{-
:m *Example Compile AST.UnMonoid
:seti -XOverloadedLabels -XGADTs
(($ SCons (Value 2) SNil) =<<) $ compile knownEnv $ fromNamed $ lambda @(TScal TF64) #x $ body $ constArr_ @TF64 (arrayGenerate (ShNil `ShCons` 10) (\(IxNil `IxCons` i) -> fromIntegral i))
-}


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

-- TODO: array lifetimes in C?


compile :: SList STy env -> Ex env t
        -> IO (SList Value env -> IO (Rep t))
compile = \env expr -> do
  let source = compileToString env expr
  hPutStrLn stderr $ "Generated C source: <<<\n" ++ source ++ ">>>"
  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)
  | 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
  | 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
  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 ((if cnst then "const " else "") ++ typ ++ " " ++ 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))
  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)
  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
  where
    precTable = Map.fromList
      [("||", (2, (2, 2)))
      ,("&&", (3, (3, 3)))
      ,("==", (4, (5, 5)))
      ,("!=", (4, (5, 5)))
      ,("<", (5, (6, 6)))
      ,(">", (5, (6, 6)))
      ,("<=", (5, (6, 6)))
      ,(">=", (5, (6, 6)))
      ,("+", (6, (6, 6)))
      ,("-", (6, (6, 7)))
      ,("*", (7, (7, 7)))
      ,("/", (7, (7, 8)))
      ,("%", (7, (7, 8)))]

repTy :: Ty -> String
repTy (TScal st) = case st of
  TI32 -> "int32_t"
  TI64 -> "int64_t"
  TF32 -> "float"
  TF64 -> "double"
  TBool -> "uint8_t"
repTy t = genStructName t

repSTy :: STy t -> String
repSTy = repTy . unSTy

genStructName :: Ty -> String
genStructName = \t -> "ty_" ++ gen t where
  -- all tags start with a letter, so the array mangling is unambiguous.
  gen :: Ty -> String
  gen TNil = "n"
  gen (TPair a b) = 'P' : gen a ++ gen b
  gen (TEither a b) = 'E' : gen a ++ gen b
  gen (TMaybe t) = 'M' : gen t
  gen (TArr n t) = "A" ++ show (fromNat n) ++ gen t
  gen (TScal st) = case st of
    TI32 -> "i"
    TI64 -> "j"
    TF32 -> "f"
    TF64 -> "d"
    TBool -> "b"
  gen (TAccum t) = 'C' : gen t

genStruct :: String -> Ty -> [StructDecl]
genStruct name topty = case topty of
  TNil ->
    [StructDecl name "" com]
  TPair a b ->
    [StructDecl name (repTy a ++ " a; " ++ repTy b ++ " b;") com]
  TEither a b ->  -- 0 -> a, 1 -> b
    [StructDecl name ("uint8_t tag; union { " ++ repTy a ++ " a; " ++ repTy b ++ " b; };") com]
  TMaybe t ->  -- 0 -> nothing, 1 -> just
    [StructDecl name ("uint8_t tag; " ++ repTy t ++ " a;") com]
  TArr n t ->
    -- The buffer is trailed by a VLA for the actual array data.
    [StructDecl (name ++ "_buf") ("size_t sh[" ++ show (fromNat n) ++ "]; size_t refc; " ++ repTy t ++ " a[];") ""
    ,StructDecl name (name ++ "_buf *buf;") com]
  TScal _ ->
    []
  TAccum t ->
    [StructDecl name (repTy t ++ " a;") com]
  where
    com = ppTy 0 topty

-- State: already-generated (skippable) struct names
-- Writer: the structs in declaration order
genStructs :: Ty -> 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
        TNil -> pure ()
        TPair a b -> genStructs a >> genStructs b
        TEither a b -> genStructs a >> genStructs b
        TMaybe t -> genStructs t
        TArr _ t -> genStructs t
        TScal _ -> pure ()
        TAccum t -> genStructs t

      tell (BList (genStruct name ty))

genAllStructs :: Foldable t => t Ty -> [StructDecl]
genAllStructs tys = toList $ evalState (execWriterT (mapM_ genStructs tys)) mempty

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

type CompM a = State CompState a

genId :: CompM Int
genId = state $ \s -> (csNextId s, s { csNextId = csNextId s + 1 })

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

genName :: CompM String
genName = genName' "x"

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

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

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

emitTLD :: String -> CompM ()
emitTLD decl = 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) = runState (compile' args expr) (CompState mempty mempty mempty 1)
      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 <stdint.h>\n"
       ,showString "#include <stdlib.h>\n"
       ,showString "#include <math.h>\n\n"
       ,compose $ map (\sd -> printStructDecl sd . showString "\n") structs
       ,showString "\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 $ map (\st -> showString "  " . printStmt 1 st . showString "\n") (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) { abort(); }\n"
       ,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"
       ,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
        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)
      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 Increment t var
    return $ CELit var

  ELet _ rhs body -> do
    e <- compile' env rhs
    var <- genName
    emit $ SVarDecl True (repSTy (typeOf rhs)) var e
    rete <- compile' (Const var `SCons` env) body
    incrementVarAlways 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 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 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"), ("a", e1)]

  EInr _ t e -> do
    name <- emitStruct (STEither t (typeOf e))
    e2 <- compile' env e
    return $ CEStruct name [("tag", CELit "1"), ("b", 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
             (BList stmts2 <> pure (SAsg retvar e2))
             (BList 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 ++ ".a") `SCons` env) a
    (e3, stmts3) <- scope $ compile' (Const (var ++ ".b") `SCons` env) b
    ((), stmtsRel1) <- scope $ incrementVarAlways Decrement t1 (var ++ ".a")
    ((), stmtsRel2) <- scope $ incrementVarAlways Decrement t2 (var ++ ".b")
    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"))
                           (BList stmts2
                            <> BList stmtsRel1
                            <> pure (SAsg retvar e2))
                           (BList stmts3
                            <> BList 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"), ("a", 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 ++ ".a") `SCons` env) b
    ((), stmtsRel) <- scope $ incrementVarAlways Decrement t (var ++ ".a")
    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"))
                           (BList stmts2
                            <> pure (SAsg retvar e2))
                           (BList stmts3
                            <> BList 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, .a = {" ++ intercalate "," (map (compileScal False t) (toList vec)) ++ "}};"
    return (CEStruct strname [("buf", CEAddrOf (CELit tldname))])

  -- EBuild _ n a b -> error "TODO" -- genStruct (STArr n (typeOf b)) <> EBuild ext n (compile' a) (compile' b)

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

  -- ESum1Inner _ e -> error "TODO" -- ESum1Inner ext (compile' e)

  -- EUnit _ e -> error "TODO" -- EUnit ext (compile' e)

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

  -- EMaximum1Inner _ e -> error "TODO" -- EMaximum1Inner ext (compile' e)

  -- EMinimum1Inner _ e -> error "TODO" -- EMinimum1Inner ext (compile' e)

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

  -- EIdx0 _ e -> error "TODO" -- EIdx0 ext (compile' e)

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

  EIdx _ earr eidx -> do
    let STArr n t = typeOf earr
    arrname <- genName
    idxname <- genName
    emit . SVarDecl True (repSTy (typeOf earr)) arrname =<< compile' env earr
    emit . SVarDecl True (repSTy (typeOf eidx)) idxname =<< compile' env eidx
    resname <- genName
    emit $ SVarDecl True (repSTy t) resname (CEIndex (CELit (arrname ++ ".buf->a")) (toLinearIdx n arrname idxname))
    incrementVarAlways Decrement (STArr n t) arrname
    return (CELit resname)

  -- EShape _ e -> error "TODO" -- EShape ext (compile' e)

  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 _ t1 t2 t3 a b c e1 e2 -> error "TODO" -- ECustom ext t1 t2 t3 (compile' a) (compile' b) (compile' c) (compile' e1) (compile' e2)

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

  -- EAccum _ n a b e -> error "TODO" -- EAccum n (compile' a) (compile' b) (compile' e)

  EError _ t s -> do
    name <- emitStruct t
    -- using 'show' here is wrong, but it's good enough for me.
    emit $ SVerbatim $ "fprintf(stderr, \"ERROR: %s\\n\", " ++ show s ++ "); exit(1);"
    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"

  _ -> error "Compile: not implemented"

data Increment = Increment | Decrement
  deriving (Show)

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

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

data ArrayTree = ATArray  -- ^ we've arrived at an array we need to decrement the refcount of
               | 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 "a" (makeArrayTree a))
                                              (smartATProj "b" (makeArrayTree b))
makeArrayTree (STMaybe t) = smartATCondTag ATNoop (makeArrayTree t)
makeArrayTree (STArr _ _) = ATArray
makeArrayTree (STScal _) = ATNoop
makeArrayTree (STAccum _) = ATNoop

incrementVar' :: Increment -> String -> ArrayTree -> CompM ()
incrementVar' inc path ATArray =
  case inc of
    Increment -> emit $ SVerbatim (path ++ ".buf->refc++;")
    Decrement ->
      emit $ SVerbatim $ "if (--" ++ path ++ ".buf->refc == 0) free(" ++ path ++ ".buf);"
incrementVar' _ _ ATNoop = pure ()
incrementVar' inc path (ATProj field t) = incrementVar' inc (path ++ "." ++ field) t
incrementVar' inc path (ATCondTag t1 t2) = do
  ((), stmts1) <- scope $ incrementVar' inc path t1
  ((), stmts2) <- scope $ incrementVar' inc path t2
  emit $ SIf (CEBinop (CELit (path ++ ".tag")) "==" (CELit "0")) (BList stmts1) (BList stmts2)
incrementVar' inc path (ATBoth t1 t2) = incrementVar' inc path t1 >> incrementVar' 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) ++ "]")))
--   _

compileOpGeneral :: SOp a b -> CExpr -> CompM CExpr
compileOpGeneral op e1 = do
  let unary cop = return @(State CompState) $ 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 @(State CompState) $ 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"

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

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