-- | Turn Core primitives into Flat - removing the folds
-- The input statements must be in A-normal form.
# LANGUAGE NoImplicitPrelude    #
# LANGUAGE PatternGuards        #
{-# LANGUAGE RecordWildCards     #-}
# LANGUAGE ScopedTypeVariables #
{-# LANGUAGE OverloadedStrings   #-}
module Icicle.Avalanche.Statement.Flatten.Exp (
    flatX, flatXS
  ) where

import              Icicle.Avalanche.Statement.Flatten.Algorithms
import              Icicle.Avalanche.Statement.Flatten.Base

import              Icicle.Avalanche.Statement.Statement
import qualified    Icicle.Avalanche.Prim.Flat      as Flat
import qualified    Icicle.Avalanche.Prim.Compounds as Flat

import qualified    Icicle.Core.Exp.Prim           as Core
import qualified    Icicle.Common.Exp.Prim.Minimal as Min

import              Icicle.Common.Base
import              Icicle.Common.Type
import              Icicle.Common.Exp
import              Icicle.Common.Exp.Simp.Beta
import              Icicle.Common.Fresh

import              Icicle.Internal.Pretty

import              P
import              Control.Monad.Trans.Class

import              Data.Hashable                  (Hashable)
import              Data.String                    (IsString)


flatXS  :: (Pretty n, Hashable n, Eq n, IsString n, Show n, Show a)
        => a
        -> [Exp a n Core.Prim]
        -> [Exp a n Flat.Prim]
        -> ([Exp a n Flat.Prim] -> FlatM a n)
        -> FlatM a n

flatXS _       []     ys stm = stm ys
flatXS a_fresh (x:xs) ys stm
 = flatX a_fresh x
 $ \x'
 -> flatXS a_fresh xs (ys <> [x']) stm


-- | Flatten an expression, wrapping the statement with any lets or loops or other bindings
-- The statement function takes the new expression.
flatX   :: forall a n . (Pretty n, Hashable n, Eq n, IsString n, Show n, Show a)
        => a
        -> Exp a n Core.Prim
        -> (Exp a n Flat.Prim -> FlatM a n)
        -> FlatM a n

flatX a_fresh xx stm
 = convX
 where
  -- Do a bit of simplification.
  Betas must be converted to lets even if they are not simple values :
  -- Unapplied lambdas aren't really able to be converted, since
  -- we're lifting some expressions to statements, and there is no statement-level
  -- lambda, only expression-level lambda.
  x' = beta
     $ betaToLets a_fresh xx

  -- Annotation plumbing
  flatX'           = flatX        a_fresh
  Flat.FlatOps{..} = Flat.flatOps a_fresh
  makeApps'        = makeApps a_fresh

  xVar   = XVar   a_fresh
  xPrim  = XPrim  a_fresh
  xValue = XValue a_fresh
  xApp   = XApp   a_fresh

  -- Convert the expression. We need to handle variables here, since they might
  -- refer to already flattened expressions.
  convX
   = case xx of
      XVar a n
        -> stm $ XVar a n
      _ -> convX'

  -- Convert the simplified expression.
  convX'
   = case x' of
      -- If it doesn't do anything interesting, we can just call the statement
      -- with the original expression
      XVar a n
       -> stm $ XVar a n
      XValue a vt bv
       -> stm $ XValue a vt bv

      XApp{}
       -- Primitive applications are where it gets interesting.
       -- See below
       | Just (p,xs) <- takePrimApps x'
       -> flatPrim p xs

       -- What is the function of this application?
       --
       -- It's not a primitive.
       -- It's not a lambda, since we did betaToLets above.
  It 's not a value , since that would n't 
       -- It's not a let-bound variable, since lets can't bind funs & so wouldn't typecheck.
       --
       -- Therefore, this should not happen for a valid program.
       | otherwise
       -> lift $ Left $ FlattenErrorApplicationNonPrimitive x'

      XPrim _ p
       -> flatPrim p []

  Unapplied lambda : this should not happen for a well - typed program
      XLam{}
       -> lift $ Left $ FlattenErrorBareLambda x'


      -- Convert expression lets into statement lets
      XLet _ n p q
       -> flatX' p
       $ \p'
       -> Let n p' <$> flatX' q stm


  -- Handle primitive applications.
  PrimFolds get turned into loops and whatnot
  flatPrim :: Core.Prim -> [Exp a n Core.Prim] -> FlatM a n
  flatPrim p xs
   = case p of
  Built - in functions that are supported differently in Source / Core vs. in Avalanche / Sea

      -- User arrays
      Core.PrimMinimal (Min.PrimBuiltinFun (Min.PrimBuiltinArray (Min.PrimBuiltinLength t)))
       -> primApps (Flat.PrimProject $ Flat.PrimProjectArrayLength t) xs []

      Core.PrimMinimal (Min.PrimBuiltinFun (Min.PrimBuiltinArray (Min.PrimBuiltinIndex t)))
       | [array, ix] <- xs
       -> flatX' array
       $ \array'
       -> flatX' ix
       $ \ix'
       -> do accN <- fresh
             stm' <- stm (xVar accN)

             let accT = SumT ErrorT t
             let lenX = xPrim (Flat.PrimProject $ Flat.PrimProjectArrayLength t) `makeApps'` [array']
             let zero = xValue IntT $ VInt 0
             let ge   = xPrim (Flat.PrimMinimal $ Min.PrimRelation Min.PrimRelationGe IntT) `makeApps'` [ix', zero]
             let lt   = xPrim (Flat.PrimMinimal $ Min.PrimRelation Min.PrimRelationLt IntT) `makeApps'` [ix', lenX]
             let if1 a b = If a b mempty

             let unsafeIndex  = xPrim (Flat.PrimUnsafe $ Flat.PrimUnsafeArrayIndex t) `makeApps'` [array', ix']
             let unsafeWrap   = xPrim (Flat.PrimMinimal $ Min.PrimConst $ Min.PrimConstRight ErrorT t) `makeApps'` [unsafeIndex]
             let writeIf      = if1 ge $ if1 lt $ Write accN unsafeWrap

             return $ InitAccumulator
                    (Accumulator accN accT (xValue accT $ VLeft $ VError ExceptIndexOutOfBounds))
                    (writeIf <> Read accN accN accT stm')


  Implement heap sort in Avalanche so that it can be melted .
      Core.PrimMinimal (Min.PrimBuiltinFun (Min.PrimBuiltinArray (Min.PrimBuiltinSort t)))
       | [array] <- xs
       -> avalancheHeapSortArray a_fresh stm flatX t array

      -- Arithmetic and simple stuff are easy, just move it over
      Core.PrimMinimal pm
       -> primApps (Flat.PrimMinimal pm) xs []

      -- Handle folds below
      Core.PrimFold pf ta
       -> flatFold pf ta xs


      -- Map: insert value into map, or if key already exists,
      -- apply update function to existing value
      Core.PrimMap (Core.PrimMapInsertOrUpdate tk tv)
       | [upd, ins, key, map]   <- xs
       -> avalancheMapInsertUpdate a_fresh stm flatX tk tv upd ins key map
       -- Map with wrong arguments
       | otherwise
       -> lift $ Left $ FlattenErrorPrimBadArgs p xs

      -- Map : delete a value
      Core.PrimMap (Core.PrimMapDelete tk tv)
       | [key, map]   <- xs
       -> avalancheMapDelete a_fresh stm flatX tk tv key map
       -- Map with wrong arguments
       | otherwise
       -> lift $ Left $ FlattenErrorPrimBadArgs p xs

      -- Map: lookup by key
      Core.PrimMap (Core.PrimMapLookup tk tv)
       | [map, key]   <- xs
       -> avalancheMapLookup a_fresh stm flatX tk tv key map
       -- Map with wrong arguments
       | otherwise
       -> lift $ Left $ FlattenErrorPrimBadArgs p xs


      -- Map: create new empty map, for each element, etc
      Core.PrimMap (Core.PrimMapMapValues tk tv tv')
       | [upd, map]   <- xs
       -> flatX' map
       $ \map'
       -> do n'keys <- fresh
             n'vals <- fresh

             stm'   <- flatX'
                     ( xPrim (Core.PrimArray $ Core.PrimArrayMap tv tv')
                       `makeApps'` [upd, xVar n'vals] )
                     (\v -> stm (mapPack tk tv' (xVar n'keys) v))

             return  $ Let n'vals (mapVals tk tv map')
                     $ Let n'keys (mapKeys tk tv map')
                     $ stm'

       -- Map with wrong arguments
       | otherwise
       -> lift $ Left $ FlattenErrorPrimBadArgs p xs


      -- Map over array: create new empty array, for each element, etc
      Core.PrimArray (Core.PrimArrayMap ta tb)
       | [upd, arr]   <- xs
       -> flatX' arr
       $ \arr'
       -> do    accN <- fresh
                stm' <- stm (xVar accN)

                let arrT = ArrayT tb

                loop <- forI'   (arrLen ta arr')                  $ \iter
                     -> fmap    (Read accN accN arrT)             $
                        slet'   (arrIx ta arr' iter)              $ \elm
                     -> flatX'  (upd `xApp` elm)                  $ \elm'
                     -> slet'   (arrUpd tb (xVar accN) iter elm') $ \arr''
                     -> return  (Write accN arr'')

                return $ InitAccumulator
                            (Accumulator accN arrT (arrNew tb (arrLen ta arr')))
                            (loop <> Read accN accN arrT stm')


       -- Map with wrong arguments
       | otherwise
       -> lift $ Left $ FlattenErrorPrimBadArgs p xs


  LatestPush b f v - > ( BufPush ( fst b ) f , BufPush ( snd b ) v )
      Core.PrimLatest (Core.PrimLatestPush i t)
       | [buf, e]    <- xs
       -> flatX' e
       $  \e'
       -> flatX' buf
       $  \buf'
       -> do   tmpN      <- fresh
               let pushed = bufPush i t buf' e'
               stm'      <- stm (xVar tmpN)
               return $ Let tmpN pushed stm'

       | otherwise
       -> lift $ Left $ FlattenErrorPrimBadArgs p xs


  LatestRead b      - > BufRead ( snd b )
      Core.PrimLatest (Core.PrimLatestRead i t)
       | [buf] <- xs
       -> flatX' buf
       $  \buf'
       -> do   tmpN       <- fresh
               let fpRead  = bufRead i t buf'
               stm'       <- stm (xVar tmpN)
               return $ Let tmpN fpRead stm'

       | otherwise
       -> lift $ Left $ FlattenErrorPrimBadArgs p xs

  Convert arguments to a simple primitive .
  -- conv is what we've already converted
  primApps p [] conv
   = stm
   $ makeApps' (xPrim p)
   $ reverse conv

  primApps p (a:as) conv
   = flatX' a
   $ \a'
   -> primApps p as (a' : conv)

  -- Handle primitive folds
  --
  -- Bool is just an if
  flatFold Core.PrimFoldBool valT [then_, else_, pred]
   = flatX' pred
   $ \pred'
   -> do -- Fresh name for accumulator and result.
  We can use same name for acc & result variables because accumulators and variables are in different scopes
         acc <- fresh
  Compute the rest of the computation , assuming we 've stored result in variable named acc
         stm' <- stm (xVar acc)
         let vunit = xValue UnitT VUnit
         sthen <- flatX' (then_ `xApp` vunit) $ (return . Write acc)
         selse <- flatX' (else_ `xApp` vunit) $ (return . Write acc)
         -- Perform if and write result
         let if_ =  If pred' sthen selse
         -- After if, read back result from accumulator and then go do the rest of the statements
         let read_ = Read acc acc valT stm'
         return (InitAccumulator (Accumulator acc valT $ xValue valT $ defaultOfType valT) (if_ <> read_))

  -- Array fold becomes a loop
  flatFold (Core.PrimFoldArray telem) valT [k, z, arr]
   = do accN <- fresh
        stm' <- stm (xVar accN)

        -- Loop body updates accumulator with k function
        loop <-  flatX' arr                             $ \arr'
             ->  forI'   (arrLen telem  arr')           $ \iter
             ->  fmap    (Read accN accN valT)          $
                 slet'   (arrIx  telem arr' iter)       $ \elm
             ->  flatX'  (makeApps' k [xVar accN, elm]) $ \x
             ->  return (Write accN x)

  Initialise accumulator with value z , execute loop , read from accumulator
        flatX' z $ \z' ->
            return (InitAccumulator (Accumulator accN valT z')
                   (loop <> Read accN accN valT stm'))


  Fold over map . Very similar to above
  flatFold (Core.PrimFoldMap tk tv) valT [k, z, mmm]
   =  flatX' mmm
   $ \mmm'
   -> do n'keys  <- fresh
         n'vals  <- fresh
         n'zips  <- fresh

         n'kv <- fresh
         n'ac <- fresh
         let k'   = XLam a_fresh n'ac  valT
                  $ XLam a_fresh n'kv (PairT tk tv)
                  ( k `makeApps'` [ xVar n'ac
                                  , projCore False tk tv $ xVar n'kv
                                  , projCore True  tk tv $ xVar n'kv ])

         res     <- flatX'
                  ( xPrim (Core.PrimFold (Core.PrimFoldArray (PairT tk tv)) valT)
                     `makeApps'` [k', z, xVar n'zips] ) stm

         return   $ Let n'keys (mapKeys tk tv mmm')
                  $ Let n'vals (mapVals tk tv mmm')
                  $ Let n'zips (arrZip  tk tv (xVar n'keys) (xVar n'vals))
                  $ res



  Fold over an option is just " maybe " combinator .
  flatFold (Core.PrimFoldOption ta) valT [xsome, xnone, opt]
   = flatX' opt
   $ \opt'
   -> do
      acc  <- fresh
      stm' <- stm (xVar acc)
      tmp  <- fresh
      let vunit = xValue UnitT VUnit
      ssome <- flatX' (xsome `xApp` (xVar tmp)) $ (return . Write acc)
      snone <- flatX' (xnone `xApp` vunit)  $ (return . Write acc)
      let if_   = If (isSome ta opt') (Let tmp (optionGet ta opt') ssome) snone
          -- After if, read back result from accumulator and then go do the rest of the statements
          read_ = Read acc acc valT stm'
      return (InitAccumulator (Accumulator acc valT $ xValue valT $ defaultOfType valT) (if_ <> read_))

  Fold over an either
  flatFold (Core.PrimFoldSum ta tb) valT [xleft, xright, scrut]
   = flatX' scrut
      $ \scrut'
      -> do
         acc  <- fresh
         stm' <- stm (xVar acc)
         tmp  <- fresh
         tmp' <- fresh
         sleft   <- flatX' (xleft  `xApp` (xVar tmp )) $ (return . Write acc)
         sright  <- flatX' (xright `xApp` (xVar tmp')) $ (return . Write acc)

         let if_   = If (isRightF ta tb scrut')
                        (Let tmp' (right ta tb scrut') sright)
                        (Let tmp  (left  ta tb scrut') sleft)
           -- After if, read back result from accumulator and then go do the rest of the statements
             read_ = Read acc acc valT stm'
         return (InitAccumulator (Accumulator acc valT $ xValue valT $ defaultOfType valT) (if_ <> read_))


  -- None of the above cases apply, so must be bad arguments
  flatFold pf rt xs
   = lift $ Left $ FlattenErrorPrimBadArgs (Core.PrimFold pf rt) xs

  Create a fst # or snd #
  proj pmin t ta tb e
   = let pm | not t
            = Min.PrimPairFst ta tb
            | otherwise
            = Min.PrimPairSnd ta tb
     in (xPrim $ pmin $ Min.PrimPair $ pm) `xApp` e
  projCore = proj Core.PrimMinimal

  slet'      = slet a_fresh
  forI'      = forI a_fresh

 | Turn Core primitives into Flat - removing the folds
 The input statements must be in A-normal form.
# LANGUAGE RecordWildCards     #
# LANGUAGE OverloadedStrings   #
 | Flatten an expression, wrapping the statement with any lets or loops or other bindings
 The statement function takes the new expression.
 Do a bit of simplification.
 Unapplied lambdas aren't really able to be converted, since
 we're lifting some expressions to statements, and there is no statement-level
 lambda, only expression-level lambda.
 Annotation plumbing
 Convert the expression. We need to handle variables here, since they might
 refer to already flattened expressions.
 Convert the simplified expression.
 If it doesn't do anything interesting, we can just call the statement
 with the original expression
 Primitive applications are where it gets interesting.
 See below
 What is the function of this application?

 It's not a primitive.
 It's not a lambda, since we did betaToLets above.
 It's not a let-bound variable, since lets can't bind funs & so wouldn't typecheck.

 Therefore, this should not happen for a valid program.
 Convert expression lets into statement lets
 Handle primitive applications.
 User arrays
 Arithmetic and simple stuff are easy, just move it over
 Handle folds below
 Map: insert value into map, or if key already exists,
 apply update function to existing value
 Map with wrong arguments
 Map : delete a value
 Map with wrong arguments
 Map: lookup by key
 Map with wrong arguments
 Map: create new empty map, for each element, etc
 Map with wrong arguments
 Map over array: create new empty array, for each element, etc
 Map with wrong arguments
 conv is what we've already converted
 Handle primitive folds

 Bool is just an if
 Fresh name for accumulator and result.
 Perform if and write result
 After if, read back result from accumulator and then go do the rest of the statements
 Array fold becomes a loop
 Loop body updates accumulator with k function
 After if, read back result from accumulator and then go do the rest of the statements
 After if, read back result from accumulator and then go do the rest of the statements
 None of the above cases apply, so must be bad arguments
# LANGUAGE NoImplicitPrelude    #
# LANGUAGE PatternGuards        #
# LANGUAGE ScopedTypeVariables #
module Icicle.Avalanche.Statement.Flatten.Exp (
    flatX, flatXS
  ) where

import              Icicle.Avalanche.Statement.Flatten.Algorithms
import              Icicle.Avalanche.Statement.Flatten.Base

import              Icicle.Avalanche.Statement.Statement
import qualified    Icicle.Avalanche.Prim.Flat      as Flat
import qualified    Icicle.Avalanche.Prim.Compounds as Flat

import qualified    Icicle.Core.Exp.Prim           as Core
import qualified    Icicle.Common.Exp.Prim.Minimal as Min

import              Icicle.Common.Base
import              Icicle.Common.Type
import              Icicle.Common.Exp
import              Icicle.Common.Exp.Simp.Beta
import              Icicle.Common.Fresh

import              Icicle.Internal.Pretty

import              P
import              Control.Monad.Trans.Class

import              Data.Hashable                  (Hashable)
import              Data.String                    (IsString)


flatXS  :: (Pretty n, Hashable n, Eq n, IsString n, Show n, Show a)
        => a
        -> [Exp a n Core.Prim]
        -> [Exp a n Flat.Prim]
        -> ([Exp a n Flat.Prim] -> FlatM a n)
        -> FlatM a n

flatXS _       []     ys stm = stm ys
flatXS a_fresh (x:xs) ys stm
 = flatX a_fresh x
 $ \x'
 -> flatXS a_fresh xs (ys <> [x']) stm


flatX   :: forall a n . (Pretty n, Hashable n, Eq n, IsString n, Show n, Show a)
        => a
        -> Exp a n Core.Prim
        -> (Exp a n Flat.Prim -> FlatM a n)
        -> FlatM a n

flatX a_fresh xx stm
 = convX
 where
  Betas must be converted to lets even if they are not simple values :
  x' = beta
     $ betaToLets a_fresh xx

  flatX'           = flatX        a_fresh
  Flat.FlatOps{..} = Flat.flatOps a_fresh
  makeApps'        = makeApps a_fresh

  xVar   = XVar   a_fresh
  xPrim  = XPrim  a_fresh
  xValue = XValue a_fresh
  xApp   = XApp   a_fresh

  convX
   = case xx of
      XVar a n
        -> stm $ XVar a n
      _ -> convX'

  convX'
   = case x' of
      XVar a n
       -> stm $ XVar a n
      XValue a vt bv
       -> stm $ XValue a vt bv

      XApp{}
       | Just (p,xs) <- takePrimApps x'
       -> flatPrim p xs

  It 's not a value , since that would n't 
       | otherwise
       -> lift $ Left $ FlattenErrorApplicationNonPrimitive x'

      XPrim _ p
       -> flatPrim p []

  Unapplied lambda : this should not happen for a well - typed program
      XLam{}
       -> lift $ Left $ FlattenErrorBareLambda x'


      XLet _ n p q
       -> flatX' p
       $ \p'
       -> Let n p' <$> flatX' q stm


  PrimFolds get turned into loops and whatnot
  flatPrim :: Core.Prim -> [Exp a n Core.Prim] -> FlatM a n
  flatPrim p xs
   = case p of
  Built - in functions that are supported differently in Source / Core vs. in Avalanche / Sea

      Core.PrimMinimal (Min.PrimBuiltinFun (Min.PrimBuiltinArray (Min.PrimBuiltinLength t)))
       -> primApps (Flat.PrimProject $ Flat.PrimProjectArrayLength t) xs []

      Core.PrimMinimal (Min.PrimBuiltinFun (Min.PrimBuiltinArray (Min.PrimBuiltinIndex t)))
       | [array, ix] <- xs
       -> flatX' array
       $ \array'
       -> flatX' ix
       $ \ix'
       -> do accN <- fresh
             stm' <- stm (xVar accN)

             let accT = SumT ErrorT t
             let lenX = xPrim (Flat.PrimProject $ Flat.PrimProjectArrayLength t) `makeApps'` [array']
             let zero = xValue IntT $ VInt 0
             let ge   = xPrim (Flat.PrimMinimal $ Min.PrimRelation Min.PrimRelationGe IntT) `makeApps'` [ix', zero]
             let lt   = xPrim (Flat.PrimMinimal $ Min.PrimRelation Min.PrimRelationLt IntT) `makeApps'` [ix', lenX]
             let if1 a b = If a b mempty

             let unsafeIndex  = xPrim (Flat.PrimUnsafe $ Flat.PrimUnsafeArrayIndex t) `makeApps'` [array', ix']
             let unsafeWrap   = xPrim (Flat.PrimMinimal $ Min.PrimConst $ Min.PrimConstRight ErrorT t) `makeApps'` [unsafeIndex]
             let writeIf      = if1 ge $ if1 lt $ Write accN unsafeWrap

             return $ InitAccumulator
                    (Accumulator accN accT (xValue accT $ VLeft $ VError ExceptIndexOutOfBounds))
                    (writeIf <> Read accN accN accT stm')


  Implement heap sort in Avalanche so that it can be melted .
      Core.PrimMinimal (Min.PrimBuiltinFun (Min.PrimBuiltinArray (Min.PrimBuiltinSort t)))
       | [array] <- xs
       -> avalancheHeapSortArray a_fresh stm flatX t array

      Core.PrimMinimal pm
       -> primApps (Flat.PrimMinimal pm) xs []

      Core.PrimFold pf ta
       -> flatFold pf ta xs


      Core.PrimMap (Core.PrimMapInsertOrUpdate tk tv)
       | [upd, ins, key, map]   <- xs
       -> avalancheMapInsertUpdate a_fresh stm flatX tk tv upd ins key map
       | otherwise
       -> lift $ Left $ FlattenErrorPrimBadArgs p xs

      Core.PrimMap (Core.PrimMapDelete tk tv)
       | [key, map]   <- xs
       -> avalancheMapDelete a_fresh stm flatX tk tv key map
       | otherwise
       -> lift $ Left $ FlattenErrorPrimBadArgs p xs

      Core.PrimMap (Core.PrimMapLookup tk tv)
       | [map, key]   <- xs
       -> avalancheMapLookup a_fresh stm flatX tk tv key map
       | otherwise
       -> lift $ Left $ FlattenErrorPrimBadArgs p xs


      Core.PrimMap (Core.PrimMapMapValues tk tv tv')
       | [upd, map]   <- xs
       -> flatX' map
       $ \map'
       -> do n'keys <- fresh
             n'vals <- fresh

             stm'   <- flatX'
                     ( xPrim (Core.PrimArray $ Core.PrimArrayMap tv tv')
                       `makeApps'` [upd, xVar n'vals] )
                     (\v -> stm (mapPack tk tv' (xVar n'keys) v))

             return  $ Let n'vals (mapVals tk tv map')
                     $ Let n'keys (mapKeys tk tv map')
                     $ stm'

       | otherwise
       -> lift $ Left $ FlattenErrorPrimBadArgs p xs


      Core.PrimArray (Core.PrimArrayMap ta tb)
       | [upd, arr]   <- xs
       -> flatX' arr
       $ \arr'
       -> do    accN <- fresh
                stm' <- stm (xVar accN)

                let arrT = ArrayT tb

                loop <- forI'   (arrLen ta arr')                  $ \iter
                     -> fmap    (Read accN accN arrT)             $
                        slet'   (arrIx ta arr' iter)              $ \elm
                     -> flatX'  (upd `xApp` elm)                  $ \elm'
                     -> slet'   (arrUpd tb (xVar accN) iter elm') $ \arr''
                     -> return  (Write accN arr'')

                return $ InitAccumulator
                            (Accumulator accN arrT (arrNew tb (arrLen ta arr')))
                            (loop <> Read accN accN arrT stm')


       | otherwise
       -> lift $ Left $ FlattenErrorPrimBadArgs p xs


  LatestPush b f v - > ( BufPush ( fst b ) f , BufPush ( snd b ) v )
      Core.PrimLatest (Core.PrimLatestPush i t)
       | [buf, e]    <- xs
       -> flatX' e
       $  \e'
       -> flatX' buf
       $  \buf'
       -> do   tmpN      <- fresh
               let pushed = bufPush i t buf' e'
               stm'      <- stm (xVar tmpN)
               return $ Let tmpN pushed stm'

       | otherwise
       -> lift $ Left $ FlattenErrorPrimBadArgs p xs


  LatestRead b      - > BufRead ( snd b )
      Core.PrimLatest (Core.PrimLatestRead i t)
       | [buf] <- xs
       -> flatX' buf
       $  \buf'
       -> do   tmpN       <- fresh
               let fpRead  = bufRead i t buf'
               stm'       <- stm (xVar tmpN)
               return $ Let tmpN fpRead stm'

       | otherwise
       -> lift $ Left $ FlattenErrorPrimBadArgs p xs

  Convert arguments to a simple primitive .
  primApps p [] conv
   = stm
   $ makeApps' (xPrim p)
   $ reverse conv

  primApps p (a:as) conv
   = flatX' a
   $ \a'
   -> primApps p as (a' : conv)

  flatFold Core.PrimFoldBool valT [then_, else_, pred]
   = flatX' pred
   $ \pred'
  We can use same name for acc & result variables because accumulators and variables are in different scopes
         acc <- fresh
  Compute the rest of the computation , assuming we 've stored result in variable named acc
         stm' <- stm (xVar acc)
         let vunit = xValue UnitT VUnit
         sthen <- flatX' (then_ `xApp` vunit) $ (return . Write acc)
         selse <- flatX' (else_ `xApp` vunit) $ (return . Write acc)
         let if_ =  If pred' sthen selse
         let read_ = Read acc acc valT stm'
         return (InitAccumulator (Accumulator acc valT $ xValue valT $ defaultOfType valT) (if_ <> read_))

  flatFold (Core.PrimFoldArray telem) valT [k, z, arr]
   = do accN <- fresh
        stm' <- stm (xVar accN)

        loop <-  flatX' arr                             $ \arr'
             ->  forI'   (arrLen telem  arr')           $ \iter
             ->  fmap    (Read accN accN valT)          $
                 slet'   (arrIx  telem arr' iter)       $ \elm
             ->  flatX'  (makeApps' k [xVar accN, elm]) $ \x
             ->  return (Write accN x)

  Initialise accumulator with value z , execute loop , read from accumulator
        flatX' z $ \z' ->
            return (InitAccumulator (Accumulator accN valT z')
                   (loop <> Read accN accN valT stm'))


  Fold over map . Very similar to above
  flatFold (Core.PrimFoldMap tk tv) valT [k, z, mmm]
   =  flatX' mmm
   $ \mmm'
   -> do n'keys  <- fresh
         n'vals  <- fresh
         n'zips  <- fresh

         n'kv <- fresh
         n'ac <- fresh
         let k'   = XLam a_fresh n'ac  valT
                  $ XLam a_fresh n'kv (PairT tk tv)
                  ( k `makeApps'` [ xVar n'ac
                                  , projCore False tk tv $ xVar n'kv
                                  , projCore True  tk tv $ xVar n'kv ])

         res     <- flatX'
                  ( xPrim (Core.PrimFold (Core.PrimFoldArray (PairT tk tv)) valT)
                     `makeApps'` [k', z, xVar n'zips] ) stm

         return   $ Let n'keys (mapKeys tk tv mmm')
                  $ Let n'vals (mapVals tk tv mmm')
                  $ Let n'zips (arrZip  tk tv (xVar n'keys) (xVar n'vals))
                  $ res



  Fold over an option is just " maybe " combinator .
  flatFold (Core.PrimFoldOption ta) valT [xsome, xnone, opt]
   = flatX' opt
   $ \opt'
   -> do
      acc  <- fresh
      stm' <- stm (xVar acc)
      tmp  <- fresh
      let vunit = xValue UnitT VUnit
      ssome <- flatX' (xsome `xApp` (xVar tmp)) $ (return . Write acc)
      snone <- flatX' (xnone `xApp` vunit)  $ (return . Write acc)
      let if_   = If (isSome ta opt') (Let tmp (optionGet ta opt') ssome) snone
          read_ = Read acc acc valT stm'
      return (InitAccumulator (Accumulator acc valT $ xValue valT $ defaultOfType valT) (if_ <> read_))

  Fold over an either
  flatFold (Core.PrimFoldSum ta tb) valT [xleft, xright, scrut]
   = flatX' scrut
      $ \scrut'
      -> do
         acc  <- fresh
         stm' <- stm (xVar acc)
         tmp  <- fresh
         tmp' <- fresh
         sleft   <- flatX' (xleft  `xApp` (xVar tmp )) $ (return . Write acc)
         sright  <- flatX' (xright `xApp` (xVar tmp')) $ (return . Write acc)

         let if_   = If (isRightF ta tb scrut')
                        (Let tmp' (right ta tb scrut') sright)
                        (Let tmp  (left  ta tb scrut') sleft)
             read_ = Read acc acc valT stm'
         return (InitAccumulator (Accumulator acc valT $ xValue valT $ defaultOfType valT) (if_ <> read_))


  flatFold pf rt xs
   = lift $ Left $ FlattenErrorPrimBadArgs (Core.PrimFold pf rt) xs

  Create a fst # or snd #
  proj pmin t ta tb e
   = let pm | not t
            = Min.PrimPairFst ta tb
            | otherwise
            = Min.PrimPairSnd ta tb
     in (xPrim $ pmin $ Min.PrimPair $ pm) `xApp` e
  projCore = proj Core.PrimMinimal

  slet'      = slet a_fresh
  forI'      = forI a_fresh
