{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE Safe #-}
-- | Hyle-to-Hyle transformations: defn inlining (the uses-at-most-once
--   heuristic, plus flatten-everything), dead-defn purging, the @--rtl-opt@
--   optimization fixpoint (partial evaluation, zero-width purging, dedupe),
--   and hoisting clocked-extern calls into device instances. Inlining is
--   sound by the substitution property (doc/hyle.md, section 7); bound
--   names in inlined bodies are freshened to avoid capture. Defns are
--   processed callee-first (the call graph is acyclic), so chains of
--   single-use defns collapse fully.
module ReWire.Hyle.Transform (inline, inlineBy, purgeUnused, partialEval, purgeDevLets, purgeZeroWidth, dedupe, optimize, hoistInstances) where

import ReWire.Annotation (Annote)
import ReWire.Error (AstError, MonadError, failAt)
import ReWire.Fix (fixPure)
import ReWire.Hyle.Interp (evalExp, evalOp, IEnv (..))
import ReWire.Hyle.Mangle (stripFreshTag)
import ReWire.Hyle.Syntax
import ReWire.Pretty (showt)

import ReWire.BitVector (BV, nat)
import Control.Arrow ((>>>))
import Control.Monad (foldM)
import Numeric.Natural (Natural)

import qualified ReWire.BitVector as BV

import Control.Monad.State.Strict (State, StateT, evalState, evalStateT, get, put, gets, modify)
import Data.HashMap.Strict (HashMap)
import Data.HashSet (HashSet)
import Data.List (sort)
import Data.Maybe (fromMaybe)

import qualified Data.HashMap.Strict as Map
import qualified Data.HashSet        as Set
import qualified Data.Text           as T

-- | Inline definitions into their call sites: all of them when @flatten@,
--   otherwise those used at most once. Definitions referenced as extern
--   models are never inlined (the reference is by name), and neither are
--   defns marked noinline -- the pin wins over @--flatten@, so a
--   noinline defn is always findable as a module in the output. The
--   inlined-away defns are removed by the final purge; a follow-up
--   partialEval fuses the slice/concat plumbing exposed by inlining and
--   drops the argument wires that fusion leaves unused.
inline :: Bool -> Program -> Program
inline :: Bool -> Program -> Program
inline Bool
flatten p :: Program
p@(Program [Extern]
exts [Defn]
ds Device
dev) = (Program -> Program
partialEval (Program -> Program) -> (Program -> Program) -> Program -> Program
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Program -> Program
purgeDevLets (Program -> Program) -> (Program -> Program) -> Program -> Program
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Program -> Program
purgeUnused) (Program -> Program) -> Program -> Program
forall a b. (a -> b) -> a -> b
$ (Text -> Bool) -> Program -> Program
inlineBy Text -> Bool
inlinable Program
p
      where inlinable :: GId -> Bool
            inlinable :: Text -> Bool
inlinable Text
g = Bool -> Bool
not (Text
g Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` HashSet Text
models)
                       Bool -> Bool -> Bool
&& Bool -> Bool
not (Text
g Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` HashSet Text
pinned)
                       Bool -> Bool -> Bool
&& (Bool
flatten Bool -> Bool -> Bool
|| Int -> Text -> HashMap Text Int -> Int
forall k v. (Eq k, Hashable k) => v -> k -> HashMap k v -> v
Map.findWithDefault Int
0 Text
g HashMap Text Int
uses Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= (Int
1 :: Int))

            models :: HashSet GId
            models :: HashSet Text
models = [Text] -> HashSet Text
forall a. (Eq a, Hashable a) => [a] -> HashSet a
Set.fromList [ Text
g | Extern
e <- [Extern]
exts, Just Text
g <- [Extern -> Maybe Text
extModel Extern
e] ]

            pinned :: HashSet GId
            pinned :: HashSet Text
pinned = [Text] -> HashSet Text
forall a. (Eq a, Hashable a) => [a] -> HashSet a
Set.fromList [ Defn -> Text
defnName Defn
d | Defn
d <- [Defn]
ds, Defn -> Bool
defnNoInline Defn
d ]

            uses :: HashMap GId Int
            uses :: HashMap Text Int
uses = (Int -> Int -> Int) -> [(Text, Int)] -> HashMap Text Int
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> [(k, v)] -> HashMap k v
Map.fromListWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ([(Text, Int)] -> HashMap Text Int)
-> [(Text, Int)] -> HashMap Text Int
forall a b. (a -> b) -> a -> b
$ (Text -> (Text, Int)) -> [Text] -> [(Text, Int)]
forall a b. (a -> b) -> [a] -> [b]
map (, Int
1) ([Text] -> [(Text, Int)]) -> [Text] -> [(Text, Int)]
forall a b. (a -> b) -> a -> b
$ (Exp -> [Text]) -> [Exp] -> [Text]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Exp -> [Text]
expCalls
                 ([Exp] -> [Text]) -> [Exp] -> [Text]
forall a b. (a -> b) -> a -> b
$ (Defn -> Exp) -> [Defn] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Defn -> Exp
defnBody [Defn]
ds [Exp] -> [Exp] -> [Exp]
forall a. Semigroup a => a -> a -> a
<> (Stmt -> Exp) -> [Stmt] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Stmt -> Exp
stmtExp (Device -> [Stmt]
devBody Device
dev)

-- | Inline definitions satisfying the predicate into their call sites.
inlineBy :: (GId -> Bool) -> Program -> Program
inlineBy :: (Text -> Bool) -> Program -> Program
inlineBy Text -> Bool
inlinable (Program [Extern]
exts [Defn]
ds Device
dev) = Program -> Program
purgeUnused (Program -> Program) -> Program -> Program
forall a b. (a -> b) -> a -> b
$ [Extern] -> [Defn] -> Device -> Program
Program [Extern]
exts [Defn]
ds' Device
dev'
      where -- | Each defn with inlining applied, built callee-first so a
            --   body pulled in by 'beta' is already fully inlined.
            doneMap :: HashMap GId Defn
            doneMap :: HashMap Text Defn
doneMap = (HashMap Text Defn -> Defn -> HashMap Text Defn)
-> HashMap Text Defn -> [Defn] -> HashMap Text Defn
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' HashMap Text Defn -> Defn -> HashMap Text Defn
step HashMap Text Defn
forall a. Monoid a => a
mempty ([Defn] -> HashMap Text Defn) -> [Defn] -> HashMap Text Defn
forall a b. (a -> b) -> a -> b
$ [Defn] -> [Defn]
topo [Defn]
ds

            step :: HashMap GId Defn -> Defn -> HashMap GId Defn
            step :: HashMap Text Defn -> Defn -> HashMap Text Defn
step HashMap Text Defn
m Defn
d = Text -> Defn -> HashMap Text Defn -> HashMap Text Defn
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
Map.insert (Defn -> Text
defnName Defn
d) (Defn
d { defnBody = inlineExp m $ defnBody d }) HashMap Text Defn
m

            ds' :: [Defn]
            ds' :: [Defn]
ds' = [ Defn -> Maybe Defn -> Defn
forall a. a -> Maybe a -> a
fromMaybe Defn
d (Maybe Defn -> Defn) -> Maybe Defn -> Defn
forall a b. (a -> b) -> a -> b
$ Text -> HashMap Text Defn -> Maybe Defn
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup (Defn -> Text
defnName Defn
d) HashMap Text Defn
doneMap | Defn
d <- [Defn]
ds ]

            dev' :: Device
            dev' :: Device
dev' = Device
dev { devBody = map inlineStmt $ devBody dev }

            inlineStmt :: Stmt -> Stmt
            inlineStmt :: Stmt -> Stmt
inlineStmt = \ case
                  SLet Annote
an Text
x Exp
e      -> Annote -> Text -> Exp -> Stmt
SLet Annote
an Text
x (Exp -> Stmt) -> Exp -> Stmt
forall a b. (a -> b) -> a -> b
$ HashMap Text Defn -> Exp -> Exp
inlineExp HashMap Text Defn
doneMap Exp
e
                  SOutput Annote
an Text
x Exp
e   -> Annote -> Text -> Exp -> Stmt
SOutput Annote
an Text
x (Exp -> Stmt) -> Exp -> Stmt
forall a b. (a -> b) -> a -> b
$ HashMap Text Defn -> Exp -> Exp
inlineExp HashMap Text Defn
doneMap Exp
e
                  SNext Annote
an Text
x Exp
e     -> Annote -> Text -> Exp -> Stmt
SNext Annote
an Text
x (Exp -> Stmt) -> Exp -> Stmt
forall a b. (a -> b) -> a -> b
$ HashMap Text Defn -> Exp -> Exp
inlineExp HashMap Text Defn
doneMap Exp
e
                  SInstIn Annote
an Text
x Text
q Exp
e -> Annote -> Text -> Text -> Exp -> Stmt
SInstIn Annote
an Text
x Text
q (Exp -> Stmt) -> Exp -> Stmt
forall a b. (a -> b) -> a -> b
$ HashMap Text Defn -> Exp -> Exp
inlineExp HashMap Text Defn
doneMap Exp
e

            -- | Callee-first (post-order DFS over the call graph).
            topo :: [Defn] -> [Defn]
            topo :: [Defn] -> [Defn]
topo [Defn]
defs = [Defn] -> [Defn]
forall a. [a] -> [a]
reverse ([Defn] -> [Defn]) -> [Defn] -> [Defn]
forall a b. (a -> b) -> a -> b
$ (HashSet Text, [Defn]) -> [Defn]
forall a b. (a, b) -> b
snd ((HashSet Text, [Defn]) -> [Defn])
-> (HashSet Text, [Defn]) -> [Defn]
forall a b. (a -> b) -> a -> b
$ ((HashSet Text, [Defn]) -> Defn -> (HashSet Text, [Defn]))
-> (HashSet Text, [Defn]) -> [Defn] -> (HashSet Text, [Defn])
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (HashSet Text, [Defn]) -> Defn -> (HashSet Text, [Defn])
visit (HashSet Text
forall a. Monoid a => a
mempty, []) [Defn]
defs
                  where defMap :: HashMap GId Defn
                        defMap :: HashMap Text Defn
defMap = [(Text, Defn)] -> HashMap Text Defn
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
Map.fromList ([(Text, Defn)] -> HashMap Text Defn)
-> [(Text, Defn)] -> HashMap Text Defn
forall a b. (a -> b) -> a -> b
$ (Defn -> (Text, Defn)) -> [Defn] -> [(Text, Defn)]
forall a b. (a -> b) -> [a] -> [b]
map (\ Defn
d -> (Defn -> Text
defnName Defn
d, Defn
d)) [Defn]
defs

                        visit :: (HashSet GId, [Defn]) -> Defn -> (HashSet GId, [Defn])
                        visit :: (HashSet Text, [Defn]) -> Defn -> (HashSet Text, [Defn])
visit acc :: (HashSet Text, [Defn])
acc@(HashSet Text
seen, [Defn]
out) Defn
d
                              | Defn -> Text
defnName Defn
d Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` HashSet Text
seen = (HashSet Text, [Defn])
acc
                              | Bool
otherwise =
                                    let (HashSet Text
seen', [Defn]
out') = ((HashSet Text, [Defn]) -> Text -> (HashSet Text, [Defn]))
-> (HashSet Text, [Defn]) -> [Text] -> (HashSet Text, [Defn])
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (HashSet Text, [Defn]) -> Text -> (HashSet Text, [Defn])
visitName (Text -> HashSet Text -> HashSet Text
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
Set.insert (Defn -> Text
defnName Defn
d) HashSet Text
seen, [Defn]
out)
                                                      ([Text] -> (HashSet Text, [Defn]))
-> [Text] -> (HashSet Text, [Defn])
forall a b. (a -> b) -> a -> b
$ Exp -> [Text]
expCalls (Exp -> [Text]) -> Exp -> [Text]
forall a b. (a -> b) -> a -> b
$ Defn -> Exp
defnBody Defn
d
                                    in (HashSet Text
seen', Defn
d Defn -> [Defn] -> [Defn]
forall a. a -> [a] -> [a]
: [Defn]
out')

                        visitName :: (HashSet GId, [Defn]) -> GId -> (HashSet GId, [Defn])
                        visitName :: (HashSet Text, [Defn]) -> Text -> (HashSet Text, [Defn])
visitName (HashSet Text, [Defn])
acc Text
g = case Text -> HashMap Text Defn -> Maybe Defn
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup Text
g HashMap Text Defn
defMap of
                              Just Defn
d  -> (HashSet Text, [Defn]) -> Defn -> (HashSet Text, [Defn])
visit (HashSet Text, [Defn])
acc Defn
d
                              Maybe Defn
Nothing -> (HashSet Text, [Defn])
acc

            inlineExp :: HashMap GId Defn -> Exp -> Exp
            inlineExp :: HashMap Text Defn -> Exp -> Exp
inlineExp HashMap Text Defn
done Exp
e0 = State Int Exp -> Int -> Exp
forall s a. State s a -> s -> a
evalState (Exp -> State Int Exp
go Exp
e0) Int
0
                  where go :: Exp -> State Int Exp
                        go :: Exp -> State Int Exp
go = \ case
                              e :: Exp
e@Lit {}            -> Exp -> State Int Exp
forall a. a -> StateT Int Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
                              e :: Exp
e@Undef {}          -> Exp -> State Int Exp
forall a. a -> StateT Int Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
                              e :: Exp
e@Var {}            -> Exp -> State Int Exp
forall a. a -> StateT Int Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
                              Cat Annote
an Exp
e1 Exp
e2        -> Annote -> Exp -> Exp -> Exp
Cat Annote
an (Exp -> Exp -> Exp)
-> State Int Exp -> StateT Int Identity (Exp -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> State Int Exp
go Exp
e1 StateT Int Identity (Exp -> Exp) -> State Int Exp -> State Int Exp
forall a b.
StateT Int Identity (a -> b)
-> StateT Int Identity a -> StateT Int Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> State Int Exp
go Exp
e2
                              Slice Annote
an Size
i Size
k Exp
e      -> Annote -> Size -> Size -> Exp -> Exp
Slice Annote
an Size
i Size
k (Exp -> Exp) -> State Int Exp -> State Int Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> State Int Exp
go Exp
e
                              Prim Annote
an Size
sz Op
op [Exp]
es    -> Annote -> Size -> Op -> [Exp] -> Exp
Prim Annote
an Size
sz Op
op ([Exp] -> Exp) -> StateT Int Identity [Exp] -> State Int Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> State Int Exp) -> [Exp] -> StateT Int Identity [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Exp -> State Int Exp
go [Exp]
es
                              XCall Annote
an Size
sz Text
x [Natural]
cs [Exp]
es -> Annote -> Size -> Text -> [Natural] -> [Exp] -> Exp
XCall Annote
an Size
sz Text
x [Natural]
cs ([Exp] -> Exp) -> StateT Int Identity [Exp] -> State Int Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> State Int Exp) -> [Exp] -> StateT Int Identity [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Exp -> State Int Exp
go [Exp]
es
                              If Annote
an Size
sz Exp
c Exp
t Exp
e      -> Annote -> Size -> Exp -> Exp -> Exp -> Exp
If Annote
an Size
sz (Exp -> Exp -> Exp -> Exp)
-> State Int Exp -> StateT Int Identity (Exp -> Exp -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> State Int Exp
go Exp
c StateT Int Identity (Exp -> Exp -> Exp)
-> State Int Exp -> StateT Int Identity (Exp -> Exp)
forall a b.
StateT Int Identity (a -> b)
-> StateT Int Identity a -> StateT Int Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> State Int Exp
go Exp
t StateT Int Identity (Exp -> Exp) -> State Int Exp -> State Int Exp
forall a b.
StateT Int Identity (a -> b)
-> StateT Int Identity a -> StateT Int Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> State Int Exp
go Exp
e
                              Let Annote
an Size
sz Text
x Exp
e1 Exp
e2   -> Annote -> Size -> Text -> Exp -> Exp -> Exp
Let Annote
an Size
sz Text
x (Exp -> Exp -> Exp)
-> State Int Exp -> StateT Int Identity (Exp -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> State Int Exp
go Exp
e1 StateT Int Identity (Exp -> Exp) -> State Int Exp -> State Int Exp
forall a b.
StateT Int Identity (a -> b)
-> StateT Int Identity a -> StateT Int Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> State Int Exp
go Exp
e2
                              Call Annote
an Size
sz Text
g [Exp]
es     -> do
                                    es' <- (Exp -> State Int Exp) -> [Exp] -> StateT Int Identity [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Exp -> State Int Exp
go [Exp]
es
                                    case Map.lookup g done of
                                          Just Defn
d | Text -> Bool
inlinable Text
g -> Annote -> Defn -> [Exp] -> State Int Exp
beta Annote
an Defn
d [Exp]
es'
                                          Maybe Defn
_                    -> Exp -> State Int Exp
forall a. a -> StateT Int Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> State Int Exp) -> Exp -> State Int Exp
forall a b. (a -> b) -> a -> b
$ Annote -> Size -> Text -> [Exp] -> Exp
Call Annote
an Size
sz Text
g [Exp]
es'

                        -- | Substitute the arguments into the (freshened)
                        --   body: atomic arguments (variables, literals)
                        --   substitute directly; the rest are bound to
                        --   freshened parameters. The body is already fully
                        --   inlined (callee-first order).
                        beta :: Annote -> Defn -> [Exp] -> State Int Exp
                        beta :: Annote -> Defn -> [Exp] -> State Int Exp
beta Annote
an (Defn Annote
_ Text
_ Sig
_ [Text]
ps Exp
body Bool
_ Blind [Text]
_) [Exp]
es = do
                              binds <- ((Text, Exp) -> StateT Int Identity ((Text, Exp), [(Text, Exp)]))
-> [(Text, Exp)]
-> StateT Int Identity [((Text, Exp), [(Text, Exp)])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Text, Exp) -> StateT Int Identity ((Text, Exp), [(Text, Exp)])
bind ([(Text, Exp)]
 -> StateT Int Identity [((Text, Exp), [(Text, Exp)])])
-> [(Text, Exp)]
-> StateT Int Identity [((Text, Exp), [(Text, Exp)])]
forall a b. (a -> b) -> a -> b
$ [Text] -> [Exp] -> [(Text, Exp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Text]
ps [Exp]
es
                              body' <- freshenExp (Map.fromList $ map fst binds) body
                              pure $ foldr (\ (Text
x, Exp
e) Exp
b -> Annote -> Size -> Text -> Exp -> Exp -> Exp
Let Annote
an (Exp -> Size
forall a. SizeAnnotated a => a -> Size
sizeOf Exp
b) Text
x Exp
e Exp
b) body' $ concatMap snd binds
                              where bind :: (Name, Exp) -> State Int ((Name, Exp), [(Name, Exp)])
                                    bind :: (Text, Exp) -> StateT Int Identity ((Text, Exp), [(Text, Exp)])
bind (Text
p, Exp
e) | Exp -> Bool
atomic Exp
e  = ((Text, Exp), [(Text, Exp)])
-> StateT Int Identity ((Text, Exp), [(Text, Exp)])
forall a. a -> StateT Int Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Text
p, Exp
e), [])
                                                | Bool
otherwise = do
                                                      p' <- Text -> State Int Text
freshName Text
p
                                                      pure ((p, Var an (sizeOf e) p'), [(p', e)])

                                    atomic :: Exp -> Bool
                                    atomic :: Exp -> Bool
atomic = \ case
                                          Var {}   -> Bool
True
                                          Lit {}   -> Bool
True
                                          Undef {} -> Bool
True
                                          Exp
_        -> Bool
False

                        -- | A fresh name for an inlined binder, seeded with
                        --   the original name (its own freshening suffixes
                        --   stripped, so suffixes don't stack across
                        --   callee-first inlining): the backends strip the
                        --   tag back off for display.
                        freshName :: Name -> State Int Name
                        freshName :: Text -> State Int Text
freshName Text
base = do
                              i <- (Int -> Int) -> StateT Int Identity Int
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Int -> Int
forall a. a -> a
id
                              modify (+ 1)
                              pure $ stripFreshTag base <> "$i" <> showt (i :: Int)

                        -- | Substitute the given free names and rename all
                        --   bound names.
                        freshenExp :: HashMap Name Exp -> Exp -> State Int Exp
                        freshenExp :: HashMap Text Exp -> Exp -> State Int Exp
freshenExp HashMap Text Exp
m = \ case
                              e :: Exp
e@Lit {}            -> Exp -> State Int Exp
forall a. a -> StateT Int Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
                              e :: Exp
e@Undef {}          -> Exp -> State Int Exp
forall a. a -> StateT Int Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
                              e :: Exp
e@(Var Annote
_ Size
_ Text
x)       -> Exp -> State Int Exp
forall a. a -> StateT Int Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> State Int Exp) -> Exp -> State Int Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Maybe Exp -> Exp
forall a. a -> Maybe a -> a
fromMaybe Exp
e (Maybe Exp -> Exp) -> Maybe Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Text -> HashMap Text Exp -> Maybe Exp
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup Text
x HashMap Text Exp
m
                              Cat Annote
an Exp
e1 Exp
e2        -> Annote -> Exp -> Exp -> Exp
Cat Annote
an (Exp -> Exp -> Exp)
-> State Int Exp -> StateT Int Identity (Exp -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashMap Text Exp -> Exp -> State Int Exp
freshenExp HashMap Text Exp
m Exp
e1 StateT Int Identity (Exp -> Exp) -> State Int Exp -> State Int Exp
forall a b.
StateT Int Identity (a -> b)
-> StateT Int Identity a -> StateT Int Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> HashMap Text Exp -> Exp -> State Int Exp
freshenExp HashMap Text Exp
m Exp
e2
                              Slice Annote
an Size
i Size
k Exp
e      -> Annote -> Size -> Size -> Exp -> Exp
Slice Annote
an Size
i Size
k (Exp -> Exp) -> State Int Exp -> State Int Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashMap Text Exp -> Exp -> State Int Exp
freshenExp HashMap Text Exp
m Exp
e
                              Prim Annote
an Size
sz Op
op [Exp]
es    -> Annote -> Size -> Op -> [Exp] -> Exp
Prim Annote
an Size
sz Op
op ([Exp] -> Exp) -> StateT Int Identity [Exp] -> State Int Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> State Int Exp) -> [Exp] -> StateT Int Identity [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (HashMap Text Exp -> Exp -> State Int Exp
freshenExp HashMap Text Exp
m) [Exp]
es
                              XCall Annote
an Size
sz Text
x [Natural]
cs [Exp]
es -> Annote -> Size -> Text -> [Natural] -> [Exp] -> Exp
XCall Annote
an Size
sz Text
x [Natural]
cs ([Exp] -> Exp) -> StateT Int Identity [Exp] -> State Int Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> State Int Exp) -> [Exp] -> StateT Int Identity [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (HashMap Text Exp -> Exp -> State Int Exp
freshenExp HashMap Text Exp
m) [Exp]
es
                              If Annote
an Size
sz Exp
c Exp
t Exp
e      -> Annote -> Size -> Exp -> Exp -> Exp -> Exp
If Annote
an Size
sz (Exp -> Exp -> Exp -> Exp)
-> State Int Exp -> StateT Int Identity (Exp -> Exp -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashMap Text Exp -> Exp -> State Int Exp
freshenExp HashMap Text Exp
m Exp
c StateT Int Identity (Exp -> Exp -> Exp)
-> State Int Exp -> StateT Int Identity (Exp -> Exp)
forall a b.
StateT Int Identity (a -> b)
-> StateT Int Identity a -> StateT Int Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> HashMap Text Exp -> Exp -> State Int Exp
freshenExp HashMap Text Exp
m Exp
t StateT Int Identity (Exp -> Exp) -> State Int Exp -> State Int Exp
forall a b.
StateT Int Identity (a -> b)
-> StateT Int Identity a -> StateT Int Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> HashMap Text Exp -> Exp -> State Int Exp
freshenExp HashMap Text Exp
m Exp
e
                              Call Annote
an Size
sz Text
g [Exp]
es     -> Annote -> Size -> Text -> [Exp] -> Exp
Call Annote
an Size
sz Text
g ([Exp] -> Exp) -> StateT Int Identity [Exp] -> State Int Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> State Int Exp) -> [Exp] -> StateT Int Identity [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (HashMap Text Exp -> Exp -> State Int Exp
freshenExp HashMap Text Exp
m) [Exp]
es
                              Let Annote
an Size
sz Text
x Exp
e1 Exp
e2   -> do
                                    e1' <- HashMap Text Exp -> Exp -> State Int Exp
freshenExp HashMap Text Exp
m Exp
e1
                                    x'  <- freshName x
                                    Let an sz x' e1' <$> freshenExp (Map.insert x (Var an (sizeOf e1) x') m) e2

stmtExp :: Stmt -> Exp
stmtExp :: Stmt -> Exp
stmtExp = \ case
      SLet Annote
_ Text
_ Exp
e      -> Exp
e
      SOutput Annote
_ Text
_ Exp
e   -> Exp
e
      SNext Annote
_ Text
_ Exp
e     -> Exp
e
      SInstIn Annote
_ Text
_ Text
_ Exp
e -> Exp
e

expFreeVars :: Exp -> HashSet Name
expFreeVars :: Exp -> HashSet Text
expFreeVars = \ case
      Lit {}            -> HashSet Text
forall a. Monoid a => a
mempty
      Undef {}          -> HashSet Text
forall a. Monoid a => a
mempty
      Var Annote
_ Size
_ Text
x         -> Text -> HashSet Text
forall a. Hashable a => a -> HashSet a
Set.singleton Text
x
      Cat Annote
_ Exp
e1 Exp
e2       -> Exp -> HashSet Text
expFreeVars Exp
e1 HashSet Text -> HashSet Text -> HashSet Text
forall a. Semigroup a => a -> a -> a
<> Exp -> HashSet Text
expFreeVars Exp
e2
      Slice Annote
_ Size
_ Size
_ Exp
e     -> Exp -> HashSet Text
expFreeVars Exp
e
      Prim Annote
_ Size
_ Op
_ [Exp]
es     -> (Exp -> HashSet Text) -> [Exp] -> HashSet Text
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Exp -> HashSet Text
expFreeVars [Exp]
es
      XCall Annote
_ Size
_ Text
_ [Natural]
_ [Exp]
es  -> (Exp -> HashSet Text) -> [Exp] -> HashSet Text
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Exp -> HashSet Text
expFreeVars [Exp]
es
      If Annote
_ Size
_ Exp
c Exp
t Exp
e      -> Exp -> HashSet Text
expFreeVars Exp
c HashSet Text -> HashSet Text -> HashSet Text
forall a. Semigroup a => a -> a -> a
<> Exp -> HashSet Text
expFreeVars Exp
t HashSet Text -> HashSet Text -> HashSet Text
forall a. Semigroup a => a -> a -> a
<> Exp -> HashSet Text
expFreeVars Exp
e
      Let Annote
_ Size
_ Text
x Exp
e1 Exp
e2   -> Exp -> HashSet Text
expFreeVars Exp
e1 HashSet Text -> HashSet Text -> HashSet Text
forall a. Semigroup a => a -> a -> a
<> Text -> HashSet Text -> HashSet Text
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
Set.delete Text
x (Exp -> HashSet Text
expFreeVars Exp
e2)
      Call Annote
_ Size
_ Text
_ [Exp]
es     -> (Exp -> HashSet Text) -> [Exp] -> HashSet Text
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Exp -> HashSet Text
expFreeVars [Exp]
es

expCalls :: Exp -> [GId]
expCalls :: Exp -> [Text]
expCalls = \ case
      Lit {}            -> []
      Undef {}          -> []
      Var {}            -> []
      Cat Annote
_ Exp
e1 Exp
e2       -> Exp -> [Text]
expCalls Exp
e1 [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> Exp -> [Text]
expCalls Exp
e2
      Slice Annote
_ Size
_ Size
_ Exp
e     -> Exp -> [Text]
expCalls Exp
e
      Prim Annote
_ Size
_ Op
_ [Exp]
es     -> (Exp -> [Text]) -> [Exp] -> [Text]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Exp -> [Text]
expCalls [Exp]
es
      XCall Annote
_ Size
_ Text
_ [Natural]
_ [Exp]
es  -> (Exp -> [Text]) -> [Exp] -> [Text]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Exp -> [Text]
expCalls [Exp]
es
      If Annote
_ Size
_ Exp
c Exp
t Exp
e      -> Exp -> [Text]
expCalls Exp
c [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> Exp -> [Text]
expCalls Exp
t [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> Exp -> [Text]
expCalls Exp
e
      Let Annote
_ Size
_ Text
_ Exp
e1 Exp
e2   -> Exp -> [Text]
expCalls Exp
e1 [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> Exp -> [Text]
expCalls Exp
e2
      Call Annote
_ Size
_ Text
g [Exp]
es     -> Text
g Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: (Exp -> [Text]) -> [Exp] -> [Text]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Exp -> [Text]
expCalls [Exp]
es

-- | Drop defns unreachable from the device body (following calls and the
--   models of called externs) and extern decls that are never referenced.
--   Not gated on noinline: pinning affects inlining, not liveness, so dead
--   pinned defns don't accumulate.
purgeUnused :: Program -> Program
purgeUnused :: Program -> Program
purgeUnused (Program [Extern]
exts [Defn]
ds Device
dev) = [Extern] -> [Defn] -> Device -> Program
Program [Extern]
exts' [Defn]
ds' Device
dev
      where defnMap :: HashMap GId Defn
            defnMap :: HashMap Text Defn
defnMap = [(Text, Defn)] -> HashMap Text Defn
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
Map.fromList ([(Text, Defn)] -> HashMap Text Defn)
-> [(Text, Defn)] -> HashMap Text Defn
forall a b. (a -> b) -> a -> b
$ (Defn -> (Text, Defn)) -> [Defn] -> [(Text, Defn)]
forall a b. (a -> b) -> [a] -> [b]
map (\ Defn
d -> (Defn -> Text
defnName Defn
d, Defn
d)) [Defn]
ds

            extMap :: HashMap Name Extern
            extMap :: HashMap Text Extern
extMap = [(Text, Extern)] -> HashMap Text Extern
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
Map.fromList ([(Text, Extern)] -> HashMap Text Extern)
-> [(Text, Extern)] -> HashMap Text Extern
forall a b. (a -> b) -> a -> b
$ (Extern -> (Text, Extern)) -> [Extern] -> [(Text, Extern)]
forall a b. (a -> b) -> [a] -> [b]
map (\ Extern
e -> (Extern -> Text
extName Extern
e, Extern
e)) [Extern]
exts

            (HashSet Text
liveDefns, HashSet Text
liveExts) = ((HashSet Text, HashSet Text)
 -> Exp -> (HashSet Text, HashSet Text))
-> (HashSet Text, HashSet Text)
-> [Exp]
-> (HashSet Text, HashSet Text)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (HashSet Text, HashSet Text) -> Exp -> (HashSet Text, HashSet Text)
visitExp (HashSet Text
forall a. Monoid a => a
mempty, [Text] -> HashSet Text
forall a. (Eq a, Hashable a) => [a] -> HashSet a
Set.fromList [ Text
ex | Instance Annote
_ Text
_ Text
ex [Natural]
_ <- Device -> [Instance]
devInstances Device
dev ])
                  ([Exp] -> (HashSet Text, HashSet Text))
-> [Exp] -> (HashSet Text, HashSet Text)
forall a b. (a -> b) -> a -> b
$ (Stmt -> Exp) -> [Stmt] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Stmt -> Exp
stmtExp ([Stmt] -> [Exp]) -> [Stmt] -> [Exp]
forall a b. (a -> b) -> a -> b
$ Device -> [Stmt]
devBody Device
dev

            visitExp :: (HashSet GId, HashSet Name) -> Exp -> (HashSet GId, HashSet Name)
            visitExp :: (HashSet Text, HashSet Text) -> Exp -> (HashSet Text, HashSet Text)
visitExp acc :: (HashSet Text, HashSet Text)
acc@(HashSet Text
lds, HashSet Text
lxs) = \ case
                  Lit {}           -> (HashSet Text, HashSet Text)
acc
                  Undef {}         -> (HashSet Text, HashSet Text)
acc
                  Var {}           -> (HashSet Text, HashSet Text)
acc
                  Cat Annote
_ Exp
e1 Exp
e2      -> (HashSet Text, HashSet Text) -> Exp -> (HashSet Text, HashSet Text)
visitExp ((HashSet Text, HashSet Text) -> Exp -> (HashSet Text, HashSet Text)
visitExp (HashSet Text, HashSet Text)
acc Exp
e1) Exp
e2
                  Slice Annote
_ Size
_ Size
_ Exp
e    -> (HashSet Text, HashSet Text) -> Exp -> (HashSet Text, HashSet Text)
visitExp (HashSet Text, HashSet Text)
acc Exp
e
                  Prim Annote
_ Size
_ Op
_ [Exp]
es    -> ((HashSet Text, HashSet Text)
 -> Exp -> (HashSet Text, HashSet Text))
-> (HashSet Text, HashSet Text)
-> [Exp]
-> (HashSet Text, HashSet Text)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (HashSet Text, HashSet Text) -> Exp -> (HashSet Text, HashSet Text)
visitExp (HashSet Text, HashSet Text)
acc [Exp]
es
                  If Annote
_ Size
_ Exp
c Exp
t Exp
e     -> ((HashSet Text, HashSet Text)
 -> Exp -> (HashSet Text, HashSet Text))
-> (HashSet Text, HashSet Text)
-> [Exp]
-> (HashSet Text, HashSet Text)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (HashSet Text, HashSet Text) -> Exp -> (HashSet Text, HashSet Text)
visitExp (HashSet Text, HashSet Text)
acc [Exp
c, Exp
t, Exp
e]
                  Let Annote
_ Size
_ Text
_ Exp
e1 Exp
e2  -> (HashSet Text, HashSet Text) -> Exp -> (HashSet Text, HashSet Text)
visitExp ((HashSet Text, HashSet Text) -> Exp -> (HashSet Text, HashSet Text)
visitExp (HashSet Text, HashSet Text)
acc Exp
e1) Exp
e2
                  Call Annote
_ Size
_ Text
g [Exp]
es    -> ((HashSet Text, HashSet Text)
 -> Exp -> (HashSet Text, HashSet Text))
-> (HashSet Text, HashSet Text)
-> [Exp]
-> (HashSet Text, HashSet Text)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (HashSet Text, HashSet Text) -> Exp -> (HashSet Text, HashSet Text)
visitExp ((HashSet Text, HashSet Text)
-> Text -> (HashSet Text, HashSet Text)
visitDefn (HashSet Text, HashSet Text)
acc Text
g) [Exp]
es
                  XCall Annote
_ Size
_ Text
x [Natural]
_ [Exp]
es -> ((HashSet Text, HashSet Text)
 -> Exp -> (HashSet Text, HashSet Text))
-> (HashSet Text, HashSet Text)
-> [Exp]
-> (HashSet Text, HashSet Text)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (HashSet Text, HashSet Text) -> Exp -> (HashSet Text, HashSet Text)
visitExp (HashSet Text, HashSet Text)
acc' [Exp]
es
                        where acc' :: (HashSet Text, HashSet Text)
acc' = case Text -> HashMap Text Extern -> Maybe Extern
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup Text
x HashMap Text Extern
extMap Maybe Extern -> (Extern -> Maybe Text) -> Maybe Text
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Extern -> Maybe Text
extModel of
                                    Just Text
g  -> (HashSet Text, HashSet Text)
-> Text -> (HashSet Text, HashSet Text)
visitDefn (HashSet Text
lds, Text -> HashSet Text -> HashSet Text
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
Set.insert Text
x HashSet Text
lxs) Text
g
                                    Maybe Text
Nothing -> (HashSet Text
lds, Text -> HashSet Text -> HashSet Text
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
Set.insert Text
x HashSet Text
lxs)

            visitDefn :: (HashSet GId, HashSet Name) -> GId -> (HashSet GId, HashSet Name)
            visitDefn :: (HashSet Text, HashSet Text)
-> Text -> (HashSet Text, HashSet Text)
visitDefn acc :: (HashSet Text, HashSet Text)
acc@(HashSet Text
lds, HashSet Text
lxs) Text
g
                  | Text
g Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` HashSet Text
lds             = (HashSet Text, HashSet Text)
acc
                  | Just Defn
d <- Text -> HashMap Text Defn -> Maybe Defn
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup Text
g HashMap Text Defn
defnMap = (HashSet Text, HashSet Text) -> Exp -> (HashSet Text, HashSet Text)
visitExp (Text -> HashSet Text -> HashSet Text
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
Set.insert Text
g HashSet Text
lds, HashSet Text
lxs) (Exp -> (HashSet Text, HashSet Text))
-> Exp -> (HashSet Text, HashSet Text)
forall a b. (a -> b) -> a -> b
$ Defn -> Exp
defnBody Defn
d
                  | Bool
otherwise                      = (HashSet Text, HashSet Text)
acc

            ds' :: [Defn]
            ds' :: [Defn]
ds' = (Defn -> Bool) -> [Defn] -> [Defn]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` HashSet Text
liveDefns) (Text -> Bool) -> (Defn -> Text) -> Defn -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Defn -> Text
defnName) [Defn]
ds

            exts' :: [Extern]
            exts' :: [Extern]
exts' = (Extern -> Bool) -> [Extern] -> [Extern]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` HashSet Text
liveExts) (Text -> Bool) -> (Extern -> Text) -> Extern -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Extern -> Text
extName) [Extern]
exts

---
--- Optimization passes: partial evaluation, zero-width purging, dedupe.
---

-- | The standard optimization pipeline, iterated to a fixpoint (bounded by
--   the --rtl-opt level).
optimize :: Natural -> Program -> Program
optimize :: Natural -> Program -> Program
optimize Natural
n = Natural -> (Program -> Program) -> Program -> Program
forall a. Eq a => Natural -> (a -> a) -> a -> a
fixPure Natural
n ((Program -> Program) -> Program -> Program)
-> (Program -> Program) -> Program -> Program
forall a b. (a -> b) -> a -> b
$ Program -> Program
partialEval (Program -> Program) -> (Program -> Program) -> Program -> Program
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Program -> Program
purgeDevLets (Program -> Program) -> (Program -> Program) -> Program -> Program
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Program -> Program
purgeZeroWidth (Program -> Program) -> (Program -> Program) -> Program -> Program
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Program -> Program
dedupe (Program -> Program) -> (Program -> Program) -> Program -> Program
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Program -> Program
purgeUnused

-- | Constant folding and wire fusion. Primitives, muxes, slices, and
--   concatenations of literals fold via the interpreter's own evaluator;
--   calls (and extern calls with models) whose arguments are all literals
--   evaluate fully. Lets bound to atomic expressions (literals, variables,
--   don't-cares) substitute away; lets bound to pure wiring (slices and
--   concatenations of atoms) are fused through where they are sliced, so
--   slice-of-concat plumbing simplifies across named wires (both expression
--   lets and device-level wires). Unused expression lets are dropped here;
--   unused device wires are dropped by 'purgeDevLets'.
partialEval :: Program -> Program
partialEval :: Program -> Program
partialEval (Program [Extern]
exts [Defn]
ds Device
dev) = [Extern] -> [Defn] -> Device -> Program
Program [Extern]
exts ((Defn -> Defn) -> [Defn] -> [Defn]
forall a b. (a -> b) -> [a] -> [b]
map Defn -> Defn
peDefn [Defn]
ds) Device
dev'
      where env :: IEnv
            env :: IEnv
env = HashMap Text Defn -> HashMap Text Extern -> IEnv
IEnv ([(Text, Defn)] -> HashMap Text Defn
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
Map.fromList ([(Text, Defn)] -> HashMap Text Defn)
-> [(Text, Defn)] -> HashMap Text Defn
forall a b. (a -> b) -> a -> b
$ (Defn -> (Text, Defn)) -> [Defn] -> [(Text, Defn)]
forall a b. (a -> b) -> [a] -> [b]
map (\ Defn
d -> (Defn -> Text
defnName Defn
d, Defn
d)) [Defn]
ds)
                       ([(Text, Extern)] -> HashMap Text Extern
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
Map.fromList ([(Text, Extern)] -> HashMap Text Extern)
-> [(Text, Extern)] -> HashMap Text Extern
forall a b. (a -> b) -> a -> b
$ (Extern -> (Text, Extern)) -> [Extern] -> [(Text, Extern)]
forall a b. (a -> b) -> [a] -> [b]
map (\ Extern
e -> (Extern -> Text
extName Extern
e, Extern
e)) [Extern]
exts)

            peDefn :: Defn -> Defn
            peDefn :: Defn -> Defn
peDefn Defn
d = Defn
d { defnBody = pe mempty $ defnBody d }

            dev' :: Device
            dev' :: Device
dev' = Device
dev { devBody = snd $ foldl' peStmt (mempty, []) $ devBody dev }

            -- | Device wires bound to pure wiring join the environment so
            --   later statements fuse through them; the wires themselves
            --   stay (purgeDevLets drops the ones fusion leaves unused).
            peStmt :: (HashMap Name Exp, [Stmt]) -> Stmt -> (HashMap Name Exp, [Stmt])
            peStmt :: (HashMap Text Exp, [Stmt]) -> Stmt -> (HashMap Text Exp, [Stmt])
peStmt (HashMap Text Exp
binds, [Stmt]
acc) = \ case
                  SLet Annote
an Text
x Exp
e ->
                        let e' :: Exp
e' = HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds Exp
e
                        in (if Exp -> Bool
wiring Exp
e' then Text -> Exp -> HashMap Text Exp -> HashMap Text Exp
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
Map.insert Text
x Exp
e' HashMap Text Exp
binds else HashMap Text Exp
binds, [Stmt]
acc [Stmt] -> [Stmt] -> [Stmt]
forall a. Semigroup a => a -> a -> a
<> [Annote -> Text -> Exp -> Stmt
SLet Annote
an Text
x Exp
e'])
                  SOutput Annote
an Text
x Exp
e   -> (HashMap Text Exp
binds, [Stmt]
acc [Stmt] -> [Stmt] -> [Stmt]
forall a. Semigroup a => a -> a -> a
<> [Annote -> Text -> Exp -> Stmt
SOutput Annote
an Text
x (Exp -> Stmt) -> Exp -> Stmt
forall a b. (a -> b) -> a -> b
$ HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds Exp
e])
                  SNext Annote
an Text
x Exp
e     -> (HashMap Text Exp
binds, [Stmt]
acc [Stmt] -> [Stmt] -> [Stmt]
forall a. Semigroup a => a -> a -> a
<> [Annote -> Text -> Exp -> Stmt
SNext Annote
an Text
x (Exp -> Stmt) -> Exp -> Stmt
forall a b. (a -> b) -> a -> b
$ HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds Exp
e])
                  SInstIn Annote
an Text
x Text
q Exp
e -> (HashMap Text Exp
binds, [Stmt]
acc [Stmt] -> [Stmt] -> [Stmt]
forall a. Semigroup a => a -> a -> a
<> [Annote -> Text -> Text -> Exp -> Stmt
SInstIn Annote
an Text
x Text
q (Exp -> Stmt) -> Exp -> Stmt
forall a b. (a -> b) -> a -> b
$ HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds Exp
e])

            pe :: HashMap Name Exp -> Exp -> Exp
            pe :: HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds = \ case
                  e :: Exp
e@(Var Annote
_ Size
_ Text
x)
                        | Just Exp
b <- Text -> HashMap Text Exp -> Maybe Exp
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup Text
x HashMap Text Exp
binds, Exp -> Bool
atomic Exp
b -> Exp
b
                        | Bool
otherwise                              -> Exp
e
                  e :: Exp
e@Lit {}   -> Exp
e
                  e :: Exp
e@Undef {} -> Exp
e
                  Cat Annote
an Exp
e1 Exp
e2 -> Annote -> [Exp] -> Exp
mergeCat Annote
an ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ (Exp -> [Exp]) -> [Exp] -> [Exp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Exp -> [Exp]
gather [HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds Exp
e1, HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds Exp
e2]
                  Slice Annote
an Size
i Size
k Exp
e -> HashMap Text Exp -> Annote -> Size -> Size -> Exp -> Exp
peSlice HashMap Text Exp
binds Annote
an Size
i Size
k (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds Exp
e
                  Prim Annote
an Size
sz Op
op [Exp]
es ->
                        let es' :: [Exp]
es' = (Exp -> Exp) -> [Exp] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds) [Exp]
es in
                        case (Exp -> Maybe BV) -> [Exp] -> Maybe [BV]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Exp -> Maybe BV
litVal [Exp]
es' of
                              Just [BV]
bvs | Right BV
bv <- Annote -> Op -> [BV] -> Either AstError BV
forall (m :: * -> *).
MonadError AstError m =>
Annote -> Op -> [BV] -> m BV
evalOp Annote
an Op
op [BV]
bvs -> Annote -> BV -> Exp
Lit Annote
an BV
bv
                              Maybe [BV]
_ -> case (Op
op, [Exp]
es') of
                                    -- The SMT-LIB division-by-zero equations
                                    -- (doc/hyle.md, section 5.2), applicable
                                    -- whenever the divisor is a zero literal.
                                    (Op
UDiv, [Exp
_, Lit Annote
_ BV
bv]) | BV -> Integer
nat BV
bv Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 -> Annote -> BV -> Exp
Lit Annote
an (BV -> Exp) -> BV -> Exp
forall a b. (a -> b) -> a -> b
$ Int -> BV
BV.ones (Int -> BV) -> Int -> BV
forall a b. (a -> b) -> a -> b
$ Size -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Size
sz
                                    (Op
UMod, [Exp
a, Lit Annote
_ BV
bv]) | BV -> Integer
nat BV
bv Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 -> Exp
a
                                    -- Reductions of a zero-width operand
                                    -- fold to the section 5.2 n = 0
                                    -- identities (the operand cannot be
                                    -- printed by the RTL backends).
                                    (Op
RedAnd, [Exp
a]) | Exp -> Size
forall a. SizeAnnotated a => a -> Size
sizeOf Exp
a Size -> Size -> Bool
forall a. Eq a => a -> a -> Bool
== Size
0 -> Annote -> BV -> Exp
Lit Annote
an (BV -> Exp) -> BV -> Exp
forall a b. (a -> b) -> a -> b
$ Int -> BV
BV.ones Int
1
                                    (Op
RedOr , [Exp
a]) | Exp -> Size
forall a. SizeAnnotated a => a -> Size
sizeOf Exp
a Size -> Size -> Bool
forall a. Eq a => a -> a -> Bool
== Size
0 -> Annote -> BV -> Exp
Lit Annote
an (BV -> Exp) -> BV -> Exp
forall a b. (a -> b) -> a -> b
$ Int -> BV
BV.zeros Int
1
                                    (Op
RedXOr, [Exp
a]) | Exp -> Size
forall a. SizeAnnotated a => a -> Size
sizeOf Exp
a Size -> Size -> Bool
forall a. Eq a => a -> a -> Bool
== Size
0 -> Annote -> BV -> Exp
Lit Annote
an (BV -> Exp) -> BV -> Exp
forall a b. (a -> b) -> a -> b
$ Int -> BV
BV.zeros Int
1
                                    -- Boolean-mux peephole (1-bit only:
                                    -- Hyle widths are load-bearing).
                                    -- Comparing a 1-bit value to a
                                    -- 1-bit literal is the value or its
                                    -- negation; double negation cancels.
                                    (Op
Eq, [Exp
a, Lit Annote
_ BV
bv]) | Exp -> Size
forall a. SizeAnnotated a => a -> Size
sizeOf Exp
a Size -> Size -> Bool
forall a. Eq a => a -> a -> Bool
== Size
1, BV -> Int
BV.width BV
bv Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 ->
                                          if BV -> Integer
nat BV
bv Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1 then Exp
a else HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Annote -> Size -> Op -> [Exp] -> Exp
Prim Annote
an Size
1 Op
Not [Exp
a]
                                    (Op
Eq, [Lit Annote
_ BV
bv, Exp
a]) | Exp -> Size
forall a. SizeAnnotated a => a -> Size
sizeOf Exp
a Size -> Size -> Bool
forall a. Eq a => a -> a -> Bool
== Size
1, BV -> Int
BV.width BV
bv Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 ->
                                          if BV -> Integer
nat BV
bv Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1 then Exp
a else HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Annote -> Size -> Op -> [Exp] -> Exp
Prim Annote
an Size
1 Op
Not [Exp
a]
                                    (Op
Not, [Prim Annote
_ Size
_ Op
Not [Exp
a]])           -> Exp
a
                                    (Op, [Exp])
_                                   -> Annote -> Size -> Op -> [Exp] -> Exp
Prim Annote
an Size
sz Op
op [Exp]
es'
                  Call Annote
an Size
sz Text
g [Exp]
es ->
                        let es' :: [Exp]
es' = (Exp -> Exp) -> [Exp] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds) [Exp]
es in
                        case (Exp -> Maybe BV) -> [Exp] -> Maybe [BV]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Exp -> Maybe BV
litVal [Exp]
es' of
                              Just [BV]
_ | Right BV
bv <- IEnv -> HashMap Text BV -> Exp -> Either AstError BV
forall (m :: * -> *).
MonadError AstError m =>
IEnv -> HashMap Text BV -> Exp -> m BV
evalExp IEnv
env HashMap Text BV
forall a. Monoid a => a
mempty (Annote -> Size -> Text -> [Exp] -> Exp
Call Annote
an Size
sz Text
g [Exp]
es') -> Annote -> BV -> Exp
Lit Annote
an BV
bv
                              Maybe [BV]
_                                                          -> Annote -> Size -> Text -> [Exp] -> Exp
Call Annote
an Size
sz Text
g [Exp]
es'
                  XCall Annote
an Size
sz Text
x [Natural]
cs [Exp]
es ->
                        let es' :: [Exp]
es' = (Exp -> Exp) -> [Exp] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds) [Exp]
es in
                        case (Exp -> Maybe BV) -> [Exp] -> Maybe [BV]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Exp -> Maybe BV
litVal [Exp]
es' of
                              Just [BV]
_ | Just Extern
ex <- Text -> HashMap Text Extern -> Maybe Extern
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup Text
x (HashMap Text Extern -> Maybe Extern)
-> HashMap Text Extern -> Maybe Extern
forall a b. (a -> b) -> a -> b
$ IEnv -> HashMap Text Extern
envExterns IEnv
env
                                     , Just Text
_  <- Extern -> Maybe Text
extModel Extern
ex
                                     , Right BV
bv <- IEnv -> HashMap Text BV -> Exp -> Either AstError BV
forall (m :: * -> *).
MonadError AstError m =>
IEnv -> HashMap Text BV -> Exp -> m BV
evalExp IEnv
env HashMap Text BV
forall a. Monoid a => a
mempty (Annote -> Size -> Text -> [Natural] -> [Exp] -> Exp
XCall Annote
an Size
sz Text
x [Natural]
cs [Exp]
es') -> Annote -> BV -> Exp
Lit Annote
an BV
bv
                              Maybe [BV]
_                                                              -> Annote -> Size -> Text -> [Natural] -> [Exp] -> Exp
XCall Annote
an Size
sz Text
x [Natural]
cs [Exp]
es'
                  If Annote
an Size
sz Exp
c Exp
t Exp
e -> case (HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds Exp
c, HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds Exp
t, HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds Exp
e) of
                        (Lit Annote
_ BV
bv, Exp
t', Exp
e') -> if BV -> Integer
nat BV
bv Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
0 then Exp
t' else Exp
e'
                        -- Boolean-mux peephole: a 1-bit mux between the
                        -- two 1-bit literals is the condition or its
                        -- negation (which re-enters 'pe' so double
                        -- negations cancel).
                        (Exp
c', Exp
t', Exp
e')
                              | Size
sz Size -> Size -> Bool
forall a. Eq a => a -> a -> Bool
== Size
1, Exp -> Size
forall a. SizeAnnotated a => a -> Size
sizeOf Exp
c' Size -> Size -> Bool
forall a. Eq a => a -> a -> Bool
== Size
1
                              , Just BV
tv <- Exp -> Maybe BV
litVal Exp
t', Just BV
ev <- Exp -> Maybe BV
litVal Exp
e'
                              , BV -> Int
BV.width BV
tv Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1, BV -> Int
BV.width BV
ev Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1, BV -> Integer
nat BV
tv Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= BV -> Integer
nat BV
ev ->
                                    if BV -> Integer
nat BV
tv Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1 then Exp
c' else HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Annote -> Size -> Op -> [Exp] -> Exp
Prim Annote
an Size
1 Op
Not [Exp
c']
                              | Bool
otherwise -> Annote -> Size -> Exp -> Exp -> Exp -> Exp
If Annote
an Size
sz Exp
c' Exp
t' Exp
e'
                  Let Annote
an Size
sz Text
x Exp
e1 Exp
e2 ->
                        let e1' :: Exp
e1'    = HashMap Text Exp -> Exp -> Exp
pe HashMap Text Exp
binds Exp
e1
                            -- Invalidate bindings an inner shadowing let
                            -- would capture (shadowing can only come from
                            -- hand-written input).
                            binds' :: HashMap Text Exp
binds' = (Exp -> Bool) -> HashMap Text Exp -> HashMap Text Exp
forall v k. (v -> Bool) -> HashMap k v -> HashMap k v
Map.filter (Bool -> Bool
not (Bool -> Bool) -> (Exp -> Bool) -> Exp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
Set.member Text
x (HashSet Text -> Bool) -> (Exp -> HashSet Text) -> Exp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> HashSet Text
expFreeVars) (HashMap Text Exp -> HashMap Text Exp)
-> HashMap Text Exp -> HashMap Text Exp
forall a b. (a -> b) -> a -> b
$ Text -> HashMap Text Exp -> HashMap Text Exp
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
Map.delete Text
x HashMap Text Exp
binds
                            e2' :: Exp
e2'    = HashMap Text Exp -> Exp -> Exp
pe (if Exp -> Bool
wiring Exp
e1' then Text -> Exp -> HashMap Text Exp -> HashMap Text Exp
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
Map.insert Text
x Exp
e1' HashMap Text Exp
binds' else HashMap Text Exp
binds') Exp
e2
                        in if Text
x Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` Exp -> HashSet Text
expFreeVars Exp
e2' then Annote -> Size -> Text -> Exp -> Exp -> Exp
Let Annote
an Size
sz Text
x Exp
e1' Exp
e2' else Exp
e2'

            litVal :: Exp -> Maybe BV
            litVal :: Exp -> Maybe BV
litVal = \ case
                  Lit Annote
_ BV
bv -> BV -> Maybe BV
forall a. a -> Maybe a
Just BV
bv
                  Exp
_        -> Maybe BV
forall a. Maybe a
Nothing

            -- | Pure wiring: evaluates to a rearrangement of its free
            --   variables' bits, with no logic. Fusing through wiring (or
            --   replicating it) duplicates no hardware.
            wiring :: Exp -> Bool
            wiring :: Exp -> Bool
wiring = \ case
                  Lit {}        -> Bool
True
                  Undef {}      -> Bool
True
                  Var {}        -> Bool
True
                  Cat Annote
_ Exp
e1 Exp
e2   -> Exp -> Bool
wiring Exp
e1 Bool -> Bool -> Bool
&& Exp -> Bool
wiring Exp
e2
                  Slice Annote
_ Size
_ Size
_ Exp
e -> Exp -> Bool
wiring Exp
e
                  Exp
_             -> Bool
False

            atomic :: Exp -> Bool
            atomic :: Exp -> Bool
atomic = \ case
                  Lit {}   -> Bool
True
                  Undef {} -> Bool
True
                  Var {}   -> Bool
True
                  Exp
_        -> Bool
False

            -- | Adjacent literals, don't-cares, and slices of the same base
            --   merge; a single piece stands alone.
            mergeCat :: Annote -> [Exp] -> Exp
            mergeCat :: Annote -> [Exp] -> Exp
mergeCat Annote
an = [Exp] -> [Exp]
go' ([Exp] -> [Exp]) -> ([Exp] -> Exp) -> [Exp] -> Exp
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> \ case
                  []  -> Annote -> BV -> Exp
Lit Annote
an BV
forall a. Monoid a => a
mempty
                  [Exp
e] -> Exp
e
                  [Exp]
es  -> (Exp -> Exp -> Exp) -> [Exp] -> Exp
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 (Annote -> Exp -> Exp -> Exp
Cat Annote
an) [Exp]
es
                  where go' :: [Exp] -> [Exp]
                        go' :: [Exp] -> [Exp]
go' = \ case
                              Lit Annote
a BV
bv : Lit Annote
_ BV
bv' : [Exp]
es   -> [Exp] -> [Exp]
go' ([Exp] -> [Exp]) -> [Exp] -> [Exp]
forall a b. (a -> b) -> a -> b
$ Annote -> BV -> Exp
Lit Annote
a (BV
bv BV -> BV -> BV
forall a. Semigroup a => a -> a -> a
<> BV
bv') Exp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
: [Exp]
es
                              Undef Annote
a Size
n : Undef Annote
_ Size
m : [Exp]
es  -> [Exp] -> [Exp]
go' ([Exp] -> [Exp]) -> [Exp] -> [Exp]
forall a b. (a -> b) -> a -> b
$ Annote -> Size -> Exp
Undef Annote
a (Size
n Size -> Size -> Size
forall a. Num a => a -> a -> a
+ Size
m) Exp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
: [Exp]
es
                              -- The left piece supplies the high bits, so
                              -- adjacency means the left slice begins where
                              -- the right one ends.
                              Slice Annote
a Size
iL Size
kL Exp
b : Slice Annote
_ Size
iR Size
kR Exp
b' : [Exp]
es
                                    | Exp
b Exp -> Exp -> Bool
forall a. Eq a => a -> a -> Bool
== Exp
b', Size
iL Size -> Size -> Bool
forall a. Eq a => a -> a -> Bool
== Size
iR Size -> Size -> Size
forall a. Num a => a -> a -> a
+ Size
kR -> [Exp] -> [Exp]
go' ([Exp] -> [Exp]) -> [Exp] -> [Exp]
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
unSlice (Annote -> Size -> Size -> Exp -> Exp
Slice Annote
a Size
iR (Size
kL Size -> Size -> Size
forall a. Num a => a -> a -> a
+ Size
kR) Exp
b) Exp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
: [Exp]
es
                              Exp
e : [Exp]
es                      -> Exp
e Exp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
: [Exp] -> [Exp]
go' [Exp]
es
                              []                          -> []

                        unSlice :: Exp -> Exp
                        unSlice :: Exp -> Exp
unSlice = \ case
                              Slice Annote
_ Size
0 Size
k Exp
e | Size
k Size -> Size -> Bool
forall a. Eq a => a -> a -> Bool
== Exp -> Size
forall a. SizeAnnotated a => a -> Size
sizeOf Exp
e -> Exp
e
                              Exp
e                             -> Exp
e

            -- | Slices of literals and don't-cares, identity slices, slices
            --   of slices, slices of concatenations (split at the piece
            --   boundaries), and slices of named wires bound to pure wiring
            --   (fused through the binding).
            peSlice :: HashMap Name Exp -> Annote -> Index -> Size -> Exp -> Exp
            peSlice :: HashMap Text Exp -> Annote -> Size -> Size -> Exp -> Exp
peSlice HashMap Text Exp
binds Annote
an Size
i Size
k Exp
e
                  | Size
k Size -> Size -> Bool
forall a. Eq a => a -> a -> Bool
== Size
0                = Annote -> BV -> Exp
Lit Annote
an BV
forall a. Monoid a => a
mempty
                  | Size
i Size -> Size -> Bool
forall a. Eq a => a -> a -> Bool
== Size
0, Size
k Size -> Size -> Bool
forall a. Eq a => a -> a -> Bool
== Exp -> Size
forall a. SizeAnnotated a => a -> Size
sizeOf Exp
e = Exp
e
                  | Bool
otherwise = case Exp
e of
                        Lit Annote
_ BV
bv        -> Annote -> BV -> Exp
Lit Annote
an (BV -> Exp) -> BV -> Exp
forall a b. (a -> b) -> a -> b
$ BV -> Size -> Size -> BV
subBV BV
bv Size
i Size
k
                        Undef Annote
_ Size
_       -> Annote -> Size -> Exp
Undef Annote
an Size
k
                        Slice Annote
_ Size
i' Size
_ Exp
e' -> HashMap Text Exp -> Annote -> Size -> Size -> Exp -> Exp
peSlice HashMap Text Exp
binds Annote
an (Size
i Size -> Size -> Size
forall a. Num a => a -> a -> a
+ Size
i') Size
k Exp
e'
                        Var Annote
_ Size
_ Text
x | Just Exp
b <- Text -> HashMap Text Exp -> Maybe Exp
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup Text
x HashMap Text Exp
binds, Bool -> Bool
not (Exp -> Bool
atomic Exp
b)
                                        -> HashMap Text Exp -> Annote -> Size -> Size -> Exp -> Exp
peSlice HashMap Text Exp
binds Annote
an Size
i Size
k Exp
b
                        Cat {}          -> Annote -> [Exp] -> Exp
mergeCat Annote
an ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ ((Exp, Size) -> Exp) -> [(Exp, Size)] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Exp, Size) -> Exp
subPiece ([(Exp, Size)] -> [Exp]) -> [(Exp, Size)] -> [Exp]
forall a b. (a -> b) -> a -> b
$ ((Exp, Size) -> Bool) -> [(Exp, Size)] -> [(Exp, Size)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Exp, Size) -> Bool
overlaps [(Exp, Size)]
pieces
                              where pieces :: [(Exp, Index)] -- MSB-first, with LSB offsets
                                    pieces :: [(Exp, Size)]
pieces = (Size, [(Exp, Size)]) -> [(Exp, Size)]
forall a b. (a, b) -> b
snd ((Size, [(Exp, Size)]) -> [(Exp, Size)])
-> (Size, [(Exp, Size)]) -> [(Exp, Size)]
forall a b. (a -> b) -> a -> b
$ (Exp -> (Size, [(Exp, Size)]) -> (Size, [(Exp, Size)]))
-> (Size, [(Exp, Size)]) -> [Exp] -> (Size, [(Exp, Size)])
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\ Exp
e' (Size
o, [(Exp, Size)]
acc) -> (Size
o Size -> Size -> Size
forall a. Num a => a -> a -> a
+ Exp -> Size
forall a. SizeAnnotated a => a -> Size
sizeOf Exp
e', (Exp
e', Size
o) (Exp, Size) -> [(Exp, Size)] -> [(Exp, Size)]
forall a. a -> [a] -> [a]
: [(Exp, Size)]
acc)) (Size
0, []) ([Exp] -> (Size, [(Exp, Size)])) -> [Exp] -> (Size, [(Exp, Size)])
forall a b. (a -> b) -> a -> b
$ Exp -> [Exp]
gather Exp
e

                                    overlaps :: (Exp, Index) -> Bool
                                    overlaps :: (Exp, Size) -> Bool
overlaps (Exp
e', Size
off) = Size
off Size -> Size -> Bool
forall a. Ord a => a -> a -> Bool
< Size
i Size -> Size -> Size
forall a. Num a => a -> a -> a
+ Size
k Bool -> Bool -> Bool
&& Size
i Size -> Size -> Bool
forall a. Ord a => a -> a -> Bool
< Size
off Size -> Size -> Size
forall a. Num a => a -> a -> a
+ Exp -> Size
forall a. SizeAnnotated a => a -> Size
sizeOf Exp
e'

                                    subPiece :: (Exp, Index) -> Exp
                                    subPiece :: (Exp, Size) -> Exp
subPiece (Exp
e', Size
off) = HashMap Text Exp -> Annote -> Size -> Size -> Exp -> Exp
peSlice HashMap Text Exp
binds Annote
an (Size -> Size -> Size
forall a. Ord a => a -> a -> a
max Size
i Size
off Size -> Size -> Size
forall a. Num a => a -> a -> a
- Size
off) (Size -> Size -> Size
forall a. Ord a => a -> a -> a
min (Size
i Size -> Size -> Size
forall a. Num a => a -> a -> a
+ Size
k) (Size
off Size -> Size -> Size
forall a. Num a => a -> a -> a
+ Exp -> Size
forall a. SizeAnnotated a => a -> Size
sizeOf Exp
e') Size -> Size -> Size
forall a. Num a => a -> a -> a
- Size -> Size -> Size
forall a. Ord a => a -> a -> a
max Size
i Size
off) Exp
e'
                        Exp
_               -> Annote -> Size -> Size -> Exp -> Exp
Slice Annote
an Size
i Size
k Exp
e

            subBV :: BV -> Index -> Size -> BV
            subBV :: BV -> Size -> Size -> BV
subBV BV
bv Size
i Size
k = Int -> Integer -> BV
forall a. Integral a => Int -> a -> BV
BV.bitVec (Size -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Size
k) (Integer -> BV) -> Integer -> BV
forall a b. (a -> b) -> a -> b
$ BV -> Integer
nat BV
bv Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` (Integer
2 Integer -> Integer -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Size -> Integer
forall a. Integral a => a -> Integer
toInteger Size
i)

-- | Drop device wires (SLets) never read by a later statement (fusion in
--   'partialEval' is what strands them).
purgeDevLets :: Program -> Program
purgeDevLets :: Program -> Program
purgeDevLets (Program [Extern]
exts [Defn]
ds Device
dev) = [Extern] -> [Defn] -> Device -> Program
Program [Extern]
exts [Defn]
ds Device
dev'
      where dev' :: Device
            dev' :: Device
dev' = Device
dev { devBody = snd $ foldr keep (mempty, []) $ devBody dev }

            keep :: Stmt -> (HashSet Name, [Stmt]) -> (HashSet Name, [Stmt])
            keep :: Stmt -> (HashSet Text, [Stmt]) -> (HashSet Text, [Stmt])
keep Stmt
s (HashSet Text
used, [Stmt]
acc) = case Stmt
s of
                  SLet Annote
_ Text
x Exp
e | Bool -> Bool
not (Text
x Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` HashSet Text
used) -> (HashSet Text
used, [Stmt]
acc)
                             | Bool
otherwise                 -> (HashSet Text
used HashSet Text -> HashSet Text -> HashSet Text
forall a. Semigroup a => a -> a -> a
<> Exp -> HashSet Text
expFreeVars Exp
e, Stmt
s Stmt -> [Stmt] -> [Stmt]
forall a. a -> [a] -> [a]
: [Stmt]
acc)
                  Stmt
_                                      -> (HashSet Text
used HashSet Text -> HashSet Text -> HashSet Text
forall a. Semigroup a => a -> a -> a
<> Exp -> HashSet Text
expFreeVars (Stmt -> Exp
stmtExp Stmt
s), Stmt
s Stmt -> [Stmt] -> [Stmt]
forall a. a -> [a] -> [a]
: [Stmt]
acc)

-- | Drop zero-width parameters (and the corresponding arguments at call
--   sites) and zero-width-result defns (calls to them become nil).
purgeZeroWidth :: Program -> Program
purgeZeroWidth :: Program -> Program
purgeZeroWidth (Program [Extern]
exts [Defn]
ds Device
dev) = [Extern] -> [Defn] -> Device -> Program
Program [Extern]
exts ((Defn -> Defn) -> [Defn] -> [Defn]
forall a b. (a -> b) -> [a] -> [b]
map Defn -> Defn
goDefn ([Defn] -> [Defn]) -> [Defn] -> [Defn]
forall a b. (a -> b) -> a -> b
$ (Defn -> Bool) -> [Defn] -> [Defn]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Size -> Size -> Bool
forall a. Ord a => a -> a -> Bool
> Size
0) (Size -> Bool) -> (Defn -> Size) -> Defn -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Defn -> Size
forall a. SizeAnnotated a => a -> Size
sizeOf) [Defn]
ds) Device
dev'
      where sigs :: HashMap GId Sig
            sigs :: HashMap Text Sig
sigs = [(Text, Sig)] -> HashMap Text Sig
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
Map.fromList ([(Text, Sig)] -> HashMap Text Sig)
-> [(Text, Sig)] -> HashMap Text Sig
forall a b. (a -> b) -> a -> b
$ (Defn -> (Text, Sig)) -> [Defn] -> [(Text, Sig)]
forall a b. (a -> b) -> [a] -> [b]
map (\ Defn
d -> (Defn -> Text
defnName Defn
d, Defn -> Sig
defnSig Defn
d)) [Defn]
ds

            goDefn :: Defn -> Defn
            goDefn :: Defn -> Defn
goDefn (Defn Annote
an Text
g (Sig Annote
san [Size]
argSzs Size
res) [Text]
ps Exp
body Bool
ni Blind [Text]
docs) = Annote
-> Text -> Sig -> [Text] -> Exp -> Bool -> Blind [Text] -> Defn
Defn Annote
an Text
g (Annote -> [Size] -> Size -> Sig
Sig Annote
san ((Size -> Bool) -> [Size] -> [Size]
forall a. (a -> Bool) -> [a] -> [a]
filter (Size -> Size -> Bool
forall a. Ord a => a -> a -> Bool
> Size
0) [Size]
argSzs) Size
res) [Text]
ps' (Exp -> Exp
goExp Exp
body) Bool
ni Blind [Text]
docs
                  where ps' :: [Text]
ps' = [ Text
p | (Text
p, Size
sz) <- [Text] -> [Size] -> [(Text, Size)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Text]
ps [Size]
argSzs, Size
sz Size -> Size -> Bool
forall a. Ord a => a -> a -> Bool
> Size
0 ]

            dev' :: Device
            dev' :: Device
dev' = Device
dev { devBody = map goStmt $ devBody dev }

            goStmt :: Stmt -> Stmt
            goStmt :: Stmt -> Stmt
goStmt = \ case
                  SLet Annote
an Text
x Exp
e      -> Annote -> Text -> Exp -> Stmt
SLet Annote
an Text
x (Exp -> Stmt) -> Exp -> Stmt
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
goExp Exp
e
                  SOutput Annote
an Text
x Exp
e   -> Annote -> Text -> Exp -> Stmt
SOutput Annote
an Text
x (Exp -> Stmt) -> Exp -> Stmt
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
goExp Exp
e
                  SNext Annote
an Text
x Exp
e     -> Annote -> Text -> Exp -> Stmt
SNext Annote
an Text
x (Exp -> Stmt) -> Exp -> Stmt
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
goExp Exp
e
                  SInstIn Annote
an Text
x Text
q Exp
e -> Annote -> Text -> Text -> Exp -> Stmt
SInstIn Annote
an Text
x Text
q (Exp -> Stmt) -> Exp -> Stmt
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
goExp Exp
e

            goExp :: Exp -> Exp
            goExp :: Exp -> Exp
goExp = \ case
                  e :: Exp
e@Lit {}   -> Exp
e
                  e :: Exp
e@Undef {} -> Exp
e
                  e :: Exp
e@Var {}   -> Exp
e
                  Cat Annote
an Exp
e1 Exp
e2 -> Annote -> Exp -> Exp -> Exp
Cat Annote
an (Exp -> Exp
goExp Exp
e1) (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
goExp Exp
e2
                  Slice Annote
an Size
i Size
k Exp
e -> Annote -> Size -> Size -> Exp -> Exp
Slice Annote
an Size
i Size
k (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
goExp Exp
e
                  Prim Annote
an Size
sz Op
op [Exp]
es -> Annote -> Size -> Op -> [Exp] -> Exp
Prim Annote
an Size
sz Op
op ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ (Exp -> Exp) -> [Exp] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> Exp
goExp [Exp]
es
                  XCall Annote
an Size
sz Text
x [Natural]
cs [Exp]
es -> Annote -> Size -> Text -> [Natural] -> [Exp] -> Exp
XCall Annote
an Size
sz Text
x [Natural]
cs ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ (Exp -> Exp) -> [Exp] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> Exp
goExp [Exp]
es
                  If Annote
an Size
sz Exp
c Exp
t Exp
e -> Annote -> Size -> Exp -> Exp -> Exp -> Exp
If Annote
an Size
sz (Exp -> Exp
goExp Exp
c) (Exp -> Exp
goExp Exp
t) (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
goExp Exp
e
                  Let Annote
an Size
sz Text
x Exp
e1 Exp
e2 -> Annote -> Size -> Text -> Exp -> Exp -> Exp
Let Annote
an Size
sz Text
x (Exp -> Exp
goExp Exp
e1) (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
goExp Exp
e2
                  Call Annote
an Size
sz Text
g [Exp]
es
                        | Size
sz Size -> Size -> Bool
forall a. Eq a => a -> a -> Bool
== Size
0   -> Annote -> BV -> Exp
Lit Annote
an BV
forall a. Monoid a => a
mempty
                        | Bool
otherwise -> case Text -> HashMap Text Sig -> Maybe Sig
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup Text
g HashMap Text Sig
sigs of
                              Just (Sig Annote
_ [Size]
argSzs Size
_) -> Annote -> Size -> Text -> [Exp] -> Exp
Call Annote
an Size
sz Text
g [ Exp -> Exp
goExp Exp
e | (Exp
e, Size
asz) <- [Exp] -> [Size] -> [(Exp, Size)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Exp]
es [Size]
argSzs, Size
asz Size -> Size -> Bool
forall a. Ord a => a -> a -> Bool
> Size
0 ]
                              Maybe Sig
Nothing               -> Annote -> Size -> Text -> [Exp] -> Exp
Call Annote
an Size
sz Text
g ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ (Exp -> Exp) -> [Exp] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> Exp
goExp [Exp]
es

-- | Redirect calls to definitions with identical signatures and bodies to a
--   single representative. Should be followed by purgeUnused. Defns marked
--   noinline are name-pinned: they take no part in merging (neither as
--   survivors nor as redirected losers). Each survivor records the losers
--   redirected to it with an appended "also: ..." doc line.
dedupe :: Program -> Program
dedupe :: Program -> Program
dedupe (Program [Extern]
exts [Defn]
ds Device
dev) = [Extern] -> [Defn] -> Device -> Program
Program ((Extern -> Extern) -> [Extern] -> [Extern]
forall a b. (a -> b) -> [a] -> [b]
map Extern -> Extern
ddExt [Extern]
exts) ((Defn -> Defn) -> [Defn] -> [Defn]
forall a b. (a -> b) -> [a] -> [b]
map (Defn -> Defn
addAlso (Defn -> Defn) -> (Defn -> Defn) -> Defn -> Defn
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Defn -> Defn
ddDefn) [Defn]
ds) Device
dev'
      where ddDefn :: Defn -> Defn
            ddDefn :: Defn -> Defn
ddDefn Defn
d = Defn
d { defnBody = ddExp $ defnBody d }

            -- | Append one doc line to each survivor naming the losers
            --   redirected to it (sorted; always distinct from the
            --   survivor by ddMap's construction).
            addAlso :: Defn -> Defn
            addAlso :: Defn -> Defn
addAlso Defn
d = case Text -> HashMap Text [Text] -> Maybe [Text]
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup (Defn -> Text
defnName Defn
d) HashMap Text [Text]
losers of
                  Just [Text]
ls -> Defn
d { defnDoc = Blind $ unBlind (defnDoc d) <> ["also: " <> T.intercalate ", " ls] }
                  Maybe [Text]
Nothing -> Defn
d

            losers :: HashMap GId [GId]
            losers :: HashMap Text [Text]
losers = [Text] -> [Text]
forall a. Ord a => [a] -> [a]
sort ([Text] -> [Text]) -> HashMap Text [Text] -> HashMap Text [Text]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Text] -> [Text] -> [Text])
-> [(Text, [Text])] -> HashMap Text [Text]
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> [(k, v)] -> HashMap k v
Map.fromListWith [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
(<>) [ (Text
g', [Text
g]) | (Text
g, Text
g') <- HashMap Text Text -> [(Text, Text)]
forall k v. HashMap k v -> [(k, v)]
Map.toList HashMap Text Text
ddMap ]

            ddExt :: Extern -> Extern
            ddExt :: Extern -> Extern
ddExt Extern
e = Extern
e { extModel = (\ Text
g -> Text -> Text -> HashMap Text Text -> Text
forall k v. (Eq k, Hashable k) => v -> k -> HashMap k v -> v
Map.findWithDefault Text
g Text
g HashMap Text Text
ddMap) <$> extModel e }

            dev' :: Device
            dev' :: Device
dev' = Device
dev { devBody = map ddStmt $ devBody dev }

            ddStmt :: Stmt -> Stmt
            ddStmt :: Stmt -> Stmt
ddStmt = \ case
                  SLet Annote
an Text
x Exp
e      -> Annote -> Text -> Exp -> Stmt
SLet Annote
an Text
x (Exp -> Stmt) -> Exp -> Stmt
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
ddExp Exp
e
                  SOutput Annote
an Text
x Exp
e   -> Annote -> Text -> Exp -> Stmt
SOutput Annote
an Text
x (Exp -> Stmt) -> Exp -> Stmt
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
ddExp Exp
e
                  SNext Annote
an Text
x Exp
e     -> Annote -> Text -> Exp -> Stmt
SNext Annote
an Text
x (Exp -> Stmt) -> Exp -> Stmt
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
ddExp Exp
e
                  SInstIn Annote
an Text
x Text
q Exp
e -> Annote -> Text -> Text -> Exp -> Stmt
SInstIn Annote
an Text
x Text
q (Exp -> Stmt) -> Exp -> Stmt
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
ddExp Exp
e

            ddExp :: Exp -> Exp
            ddExp :: Exp -> Exp
ddExp = \ case
                  e :: Exp
e@Lit {}   -> Exp
e
                  e :: Exp
e@Undef {} -> Exp
e
                  e :: Exp
e@Var {}   -> Exp
e
                  Cat Annote
an Exp
e1 Exp
e2 -> Annote -> Exp -> Exp -> Exp
Cat Annote
an (Exp -> Exp
ddExp Exp
e1) (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
ddExp Exp
e2
                  Slice Annote
an Size
i Size
k Exp
e -> Annote -> Size -> Size -> Exp -> Exp
Slice Annote
an Size
i Size
k (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
ddExp Exp
e
                  Prim Annote
an Size
sz Op
op [Exp]
es -> Annote -> Size -> Op -> [Exp] -> Exp
Prim Annote
an Size
sz Op
op ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ (Exp -> Exp) -> [Exp] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> Exp
ddExp [Exp]
es
                  XCall Annote
an Size
sz Text
x [Natural]
cs [Exp]
es -> Annote -> Size -> Text -> [Natural] -> [Exp] -> Exp
XCall Annote
an Size
sz Text
x [Natural]
cs ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ (Exp -> Exp) -> [Exp] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> Exp
ddExp [Exp]
es
                  If Annote
an Size
sz Exp
c Exp
t Exp
e -> Annote -> Size -> Exp -> Exp -> Exp -> Exp
If Annote
an Size
sz (Exp -> Exp
ddExp Exp
c) (Exp -> Exp
ddExp Exp
t) (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
ddExp Exp
e
                  Let Annote
an Size
sz Text
x Exp
e1 Exp
e2 -> Annote -> Size -> Text -> Exp -> Exp -> Exp
Let Annote
an Size
sz Text
x (Exp -> Exp
ddExp Exp
e1) (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
ddExp Exp
e2
                  Call Annote
an Size
sz Text
g [Exp]
es -> Annote -> Size -> Text -> [Exp] -> Exp
Call Annote
an Size
sz (Text -> Text -> HashMap Text Text -> Text
forall k v. (Eq k, Hashable k) => v -> k -> HashMap k v -> v
Map.findWithDefault Text
g Text
g HashMap Text Text
ddMap) ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ (Exp -> Exp) -> [Exp] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> Exp
ddExp [Exp]
es

            -- Note: Annote's Eq/Hashable are trivial, so annotations don't
            -- perturb the keys. noinline defns are excluded on both sides
            -- (they are name-pinned).
            ddMap :: HashMap GId GId
            ddMap :: HashMap Text Text
ddMap = (Defn -> HashMap Text Text -> HashMap Text Text)
-> HashMap Text Text -> [Defn] -> HashMap Text Text
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\ (Defn Annote
_ Text
g Sig
sig [Text]
ps Exp
body Bool
_ Blind [Text]
_) -> (HashMap Text Text -> HashMap Text Text)
-> (Text -> HashMap Text Text -> HashMap Text Text)
-> Maybe Text
-> HashMap Text Text
-> HashMap Text Text
forall b a. b -> (a -> b) -> Maybe a -> b
maybe HashMap Text Text -> HashMap Text Text
forall a. a -> a
id (\ Text
g' -> if Text
g' Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
/= Text
g then Text -> Text -> HashMap Text Text -> HashMap Text Text
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
Map.insert Text
g Text
g' else HashMap Text Text -> HashMap Text Text
forall a. a -> a
id)
                              (Maybe Text -> HashMap Text Text -> HashMap Text Text)
-> Maybe Text -> HashMap Text Text -> HashMap Text Text
forall a b. (a -> b) -> a -> b
$ (Sig, [Text], Exp) -> HashMap (Sig, [Text], Exp) Text -> Maybe Text
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup (Sig
sig, [Text]
ps, Exp
body) HashMap (Sig, [Text], Exp) Text
bodies) HashMap Text Text
forall a. Monoid a => a
mempty [Defn]
unpinned
                  where bodies :: HashMap (Sig, [Name], Exp) GId
                        bodies :: HashMap (Sig, [Text], Exp) Text
bodies = (Defn
 -> HashMap (Sig, [Text], Exp) Text
 -> HashMap (Sig, [Text], Exp) Text)
-> HashMap (Sig, [Text], Exp) Text
-> [Defn]
-> HashMap (Sig, [Text], Exp) Text
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\ (Defn Annote
_ Text
g Sig
sig [Text]
ps Exp
body Bool
_ Blind [Text]
_) -> (Text -> Text -> Text)
-> (Sig, [Text], Exp)
-> Text
-> HashMap (Sig, [Text], Exp) Text
-> HashMap (Sig, [Text], Exp) Text
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
Map.insertWith Text -> Text -> Text
keepBetter (Sig
sig, [Text]
ps, Exp
body) Text
g) HashMap (Sig, [Text], Exp) Text
forall a. Monoid a => a
mempty [Defn]
unpinned

                        unpinned :: [Defn]
                        unpinned :: [Defn]
unpinned = (Defn -> Bool) -> [Defn] -> [Defn]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Defn -> Bool) -> Defn -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Defn -> Bool
defnNoInline) [Defn]
ds

            -- | The surviving name for a set of identical defns: prefer
            --   user-derived names (no '$') over instance names
            --   ('iter$W8') over compiler-fresh ones ('$'-led); then
            --   shortest, then lexicographic — a total order independent
            --   of program order and machine, so the winner never flips
            --   between compiles.
            keepBetter :: GId -> GId -> GId
            keepBetter :: Text -> Text -> Text
keepBetter Text
a Text
b | Text -> (Int, Int, Text)
rank Text
a (Int, Int, Text) -> (Int, Int, Text) -> Bool
forall a. Ord a => a -> a -> Bool
<= Text -> (Int, Int, Text)
rank Text
b = Text
a
                           | Bool
otherwise        = Text
b
                  where rank :: GId -> (Int, Int, GId)
                        rank :: Text -> (Int, Int, Text)
rank Text
g = (Text -> Int
cls Text
g, Text -> Int
T.length Text
g, Text
g)

                        cls :: GId -> Int
                        cls :: Text -> Int
cls Text
g | Bool -> Bool
not (Text
"$" Text -> Text -> Bool
`T.isInfixOf` Text
g) = Int
0
                              | Text
"$" Text -> Text -> Bool
`T.isPrefixOf` Text
g      = Int
2
                              | Bool
otherwise                 = Int
1

---
--- Hoisting clocked externs into device instances.
---

-- | Defns containing (transitively) calls to sequential externs are inlined
--   into the device body; sequential-extern calls in the device body then
--   become instances: nested lets are flattened to device-level lets (sound:
--   the language is pure and total) and each seq-XCall is replaced by the
--   concatenation of a fresh instance's output ports, with its inputs driven
--   from the call's arguments.
hoistInstances :: forall m. MonadError AstError m => Program -> m Program
hoistInstances :: forall (m :: * -> *). MonadError AstError m => Program -> m Program
hoistInstances p :: Program
p@(Program [Extern]
exts [Defn]
ds Device
_)
      | HashSet Text -> Bool
forall a. HashSet a -> Bool
Set.null HashSet Text
clocked = Program -> m Program
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Program
p
      | Bool
otherwise = (StateT Int m Program -> Int -> m Program)
-> Int -> StateT Int m Program -> m Program
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT Int m Program -> Int -> m Program
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (Int
0 :: Int) (StateT Int m Program -> m Program)
-> StateT Int m Program -> m Program
forall a b. (a -> b) -> a -> b
$ do
            -- This inlineBy is a correctness transform (clocked-extern
            -- calls must reach the device body to become instances), so it
            -- deliberately ignores defnNoInline.
            let Program [Extern]
exts' [Defn]
ds' Device
dev' = (Text -> Bool) -> Program -> Program
inlineBy (Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` HashSet Text
clocked) Program
p
            (insts, stmts) <- (([Instance], [Stmt]) -> Stmt -> StateT Int m ([Instance], [Stmt]))
-> ([Instance], [Stmt])
-> [Stmt]
-> StateT Int m ([Instance], [Stmt])
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM ([Instance], [Stmt]) -> Stmt -> StateT Int m ([Instance], [Stmt])
hoistStmt ([], []) ([Stmt] -> StateT Int m ([Instance], [Stmt]))
-> [Stmt] -> StateT Int m ([Instance], [Stmt])
forall a b. (a -> b) -> a -> b
$ Device -> [Stmt]
devBody Device
dev'
            pure $ Program exts' ds' $ dev' { devInstances = reverse insts, devBody = reverse stmts }
      where externs :: HashMap Name Extern
            externs :: HashMap Text Extern
externs = [(Text, Extern)] -> HashMap Text Extern
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
Map.fromList ([(Text, Extern)] -> HashMap Text Extern)
-> [(Text, Extern)] -> HashMap Text Extern
forall a b. (a -> b) -> a -> b
$ (Extern -> (Text, Extern)) -> [Extern] -> [(Text, Extern)]
forall a b. (a -> b) -> [a] -> [b]
map (\ Extern
e -> (Extern -> Text
extName Extern
e, Extern
e)) [Extern]
exts

            isSeq :: Name -> Bool
            isSeq :: Text -> Bool
isSeq Text
x = case Extern -> ExternKind
extKind (Extern -> ExternKind) -> Maybe Extern -> Maybe ExternKind
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> HashMap Text Extern -> Maybe Extern
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup Text
x HashMap Text Extern
externs of
                  Just (Seq Maybe Text
_ Maybe Text
_) -> Bool
True
                  Maybe ExternKind
_              -> Bool
False

            -- | Defns whose bodies transitively contain seq-extern calls.
            clocked :: HashSet GId
            clocked :: HashSet Text
clocked = (HashSet Text -> HashSet Text) -> HashSet Text -> HashSet Text
forall {t}. Eq t => (t -> t) -> t -> t
fixSet HashSet Text -> HashSet Text
step HashSet Text
direct
                  where direct :: HashSet Text
direct = [Text] -> HashSet Text
forall a. (Eq a, Hashable a) => [a] -> HashSet a
Set.fromList [ Defn -> Text
defnName Defn
d | Defn
d <- [Defn]
ds, Exp -> Bool
anySeq (Exp -> Bool) -> Exp -> Bool
forall a b. (a -> b) -> a -> b
$ Defn -> Exp
defnBody Defn
d ]

                        fixSet :: (t -> t) -> t -> t
fixSet t -> t
f t
s | t
s' t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
s   = t
s
                                   | Bool
otherwise = (t -> t) -> t -> t
fixSet t -> t
f t
s'
                              where s' :: t
s' = t -> t
f t
s

                        step :: HashSet Text -> HashSet Text
step HashSet Text
s = [Text] -> HashSet Text
forall a. (Eq a, Hashable a) => [a] -> HashSet a
Set.fromList [ Defn -> Text
defnName Defn
d
                                              | Defn
d <- [Defn]
ds
                                              , Defn -> Text
defnName Defn
d Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` HashSet Text
s
                                                    Bool -> Bool -> Bool
|| (Text -> Bool) -> [Text] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` HashSet Text
s) (Exp -> [Text]
expCalls (Exp -> [Text]) -> Exp -> [Text]
forall a b. (a -> b) -> a -> b
$ Defn -> Exp
defnBody Defn
d) ]

            anySeq :: Exp -> Bool
            anySeq :: Exp -> Bool
anySeq = \ case
                  XCall Annote
_ Size
_ Text
x [Natural]
_ [Exp]
es -> Text -> Bool
isSeq Text
x Bool -> Bool -> Bool
|| (Exp -> Bool) -> [Exp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Exp -> Bool
anySeq [Exp]
es
                  Cat Annote
_ Exp
e1 Exp
e2      -> Exp -> Bool
anySeq Exp
e1 Bool -> Bool -> Bool
|| Exp -> Bool
anySeq Exp
e2
                  Slice Annote
_ Size
_ Size
_ Exp
e    -> Exp -> Bool
anySeq Exp
e
                  Prim Annote
_ Size
_ Op
_ [Exp]
es    -> (Exp -> Bool) -> [Exp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Exp -> Bool
anySeq [Exp]
es
                  Call Annote
_ Size
_ Text
_ [Exp]
es    -> (Exp -> Bool) -> [Exp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Exp -> Bool
anySeq [Exp]
es
                  If Annote
_ Size
_ Exp
c Exp
t Exp
e     -> Exp -> Bool
anySeq Exp
c Bool -> Bool -> Bool
|| Exp -> Bool
anySeq Exp
t Bool -> Bool -> Bool
|| Exp -> Bool
anySeq Exp
e
                  Let Annote
_ Size
_ Text
_ Exp
e1 Exp
e2  -> Exp -> Bool
anySeq Exp
e1 Bool -> Bool -> Bool
|| Exp -> Bool
anySeq Exp
e2
                  Exp
_                -> Bool
False

            freshHoisted :: Name -> StateT Int m Name
            freshHoisted :: Text -> StateT Int m Text
freshHoisted Text
pfx = do
                  i <- StateT Int m Int
forall s (m :: * -> *). MonadState s m => m s
get
                  put $ i + 1
                  pure $ pfx <> showt i

            -- | Flatten an expression's nested lets to device-level lets
            --   (renamed fresh, seeded from the source binder and the
            --   extern's name so the readable base survives the hoist),
            --   then convert its seq-XCalls.
            hoistStmt :: ([Instance], [Stmt]) -> Stmt -> StateT Int m ([Instance], [Stmt])
            hoistStmt :: ([Instance], [Stmt]) -> Stmt -> StateT Int m ([Instance], [Stmt])
hoistStmt ([Instance]
insts, [Stmt]
stmts) Stmt
stmt = do
                  let (Annote
an, Exp -> Stmt
rebuild, Exp
e) = Stmt -> (Annote, Exp -> Stmt, Exp)
openStmt Stmt
stmt
                  (e', insts', stmts') <- Annote
-> HashMap Text Exp
-> ([Instance], [Stmt])
-> Exp
-> StateT Int m (Exp, [Instance], [Stmt])
hoistExp Annote
an HashMap Text Exp
forall a. Monoid a => a
mempty ([Instance]
insts, [Stmt]
stmts) Exp
e
                  pure (insts', rebuild e' : stmts')

            openStmt :: Stmt -> (Annote, Exp -> Stmt, Exp)
            openStmt :: Stmt -> (Annote, Exp -> Stmt, Exp)
openStmt = \ case
                  SLet Annote
an Text
x Exp
e      -> (Annote
an, Annote -> Text -> Exp -> Stmt
SLet Annote
an Text
x, Exp
e)
                  SOutput Annote
an Text
x Exp
e   -> (Annote
an, Annote -> Text -> Exp -> Stmt
SOutput Annote
an Text
x, Exp
e)
                  SNext Annote
an Text
x Exp
e     -> (Annote
an, Annote -> Text -> Exp -> Stmt
SNext Annote
an Text
x, Exp
e)
                  SInstIn Annote
an Text
x Text
q Exp
e -> (Annote
an, Annote -> Text -> Text -> Exp -> Stmt
SInstIn Annote
an Text
x Text
q, Exp
e)

            -- | Bottom-up: substitute renamed lets, lift each Let binding to
            --   a device-level SLet, and convert seq-XCalls to instances.
            hoistExp :: Annote -> HashMap Name Exp -> ([Instance], [Stmt]) -> Exp -> StateT Int m (Exp, [Instance], [Stmt])
            hoistExp :: Annote
-> HashMap Text Exp
-> ([Instance], [Stmt])
-> Exp
-> StateT Int m (Exp, [Instance], [Stmt])
hoistExp Annote
an HashMap Text Exp
sub acc :: ([Instance], [Stmt])
acc@([Instance]
insts, [Stmt]
stmts) = \ case
                  e :: Exp
e@Lit {}   -> (Exp, [Instance], [Stmt]) -> StateT Int m (Exp, [Instance], [Stmt])
forall a. a -> StateT Int m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, [Instance]
insts, [Stmt]
stmts)
                  e :: Exp
e@Undef {} -> (Exp, [Instance], [Stmt]) -> StateT Int m (Exp, [Instance], [Stmt])
forall a. a -> StateT Int m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, [Instance]
insts, [Stmt]
stmts)
                  e :: Exp
e@(Var Annote
_ Size
_ Text
x) -> (Exp, [Instance], [Stmt]) -> StateT Int m (Exp, [Instance], [Stmt])
forall a. a -> StateT Int m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> Maybe Exp -> Exp
forall a. a -> Maybe a -> a
fromMaybe Exp
e (Maybe Exp -> Exp) -> Maybe Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Text -> HashMap Text Exp -> Maybe Exp
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup Text
x HashMap Text Exp
sub, [Instance]
insts, [Stmt]
stmts)
                  Cat Annote
an' Exp
e1 Exp
e2 -> do
                        (e1', i1, s1) <- Annote
-> HashMap Text Exp
-> ([Instance], [Stmt])
-> Exp
-> StateT Int m (Exp, [Instance], [Stmt])
hoistExp Annote
an HashMap Text Exp
sub ([Instance], [Stmt])
acc Exp
e1
                        (e2', i2, s2) <- hoistExp an sub (i1, s1) e2
                        pure (Cat an' e1' e2', i2, s2)
                  Slice Annote
an' Size
i Size
k Exp
e -> do
                        (e', i1, s1) <- Annote
-> HashMap Text Exp
-> ([Instance], [Stmt])
-> Exp
-> StateT Int m (Exp, [Instance], [Stmt])
hoistExp Annote
an HashMap Text Exp
sub ([Instance], [Stmt])
acc Exp
e
                        pure (Slice an' i k e', i1, s1)
                  Prim Annote
an' Size
sz Op
op [Exp]
es -> do
                        (es', i1, s1) <- Annote
-> HashMap Text Exp
-> ([Instance], [Stmt])
-> [Exp]
-> StateT Int m ([Exp], [Instance], [Stmt])
hoistExps Annote
an HashMap Text Exp
sub ([Instance], [Stmt])
acc [Exp]
es
                        pure (Prim an' sz op es', i1, s1)
                  Call Annote
an' Size
sz Text
g [Exp]
es -> do
                        (es', i1, s1) <- Annote
-> HashMap Text Exp
-> ([Instance], [Stmt])
-> [Exp]
-> StateT Int m ([Exp], [Instance], [Stmt])
hoistExps Annote
an HashMap Text Exp
sub ([Instance], [Stmt])
acc [Exp]
es
                        pure (Call an' sz g es', i1, s1)
                  If Annote
an' Size
sz Exp
c Exp
t Exp
e -> do
                        (c', i1, s1) <- Annote
-> HashMap Text Exp
-> ([Instance], [Stmt])
-> Exp
-> StateT Int m (Exp, [Instance], [Stmt])
hoistExp Annote
an HashMap Text Exp
sub ([Instance], [Stmt])
acc Exp
c
                        (t', i2, s2) <- hoistExp an sub (i1, s1) t
                        (e', i3, s3) <- hoistExp an sub (i2, s2) e
                        pure (If an' sz c' t' e', i3, s3)
                  Let Annote
an' Size
_ Text
x Exp
e1 Exp
e2 -> do
                        (e1', i1, s1) <- Annote
-> HashMap Text Exp
-> ([Instance], [Stmt])
-> Exp
-> StateT Int m (Exp, [Instance], [Stmt])
hoistExp Annote
an HashMap Text Exp
sub ([Instance], [Stmt])
acc Exp
e1
                        x' <- freshHoisted $ x <> "$h"
                        hoistExp an (Map.insert x (Var an' (sizeOf e1') x') sub) (i1, SLet an' x' e1' : s1) e2
                  XCall Annote
an' Size
sz Text
x [Natural]
cs [Exp]
es
                        | Text -> Bool
isSeq Text
x -> do
                              (es', i1, s1) <- Annote
-> HashMap Text Exp
-> ([Instance], [Stmt])
-> [Exp]
-> StateT Int m ([Exp], [Instance], [Stmt])
hoistExps Annote
an HashMap Text Exp
sub ([Instance], [Stmt])
acc [Exp]
es
                              i  <- freshHoisted $ x <> "$x"
                              case Map.lookup x externs of
                                    Maybe Extern
Nothing -> Annote -> Text -> StateT Int m (Exp, [Instance], [Stmt])
forall (m :: * -> *) an a.
(MonadError AstError m, Annotation an) =>
an -> Text -> m a
failAt Annote
an' (Text -> StateT Int m (Exp, [Instance], [Stmt]))
-> Text -> StateT Int m (Exp, [Instance], [Stmt])
forall a b. (a -> b) -> a -> b
$ Text
"unknown extern: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
x
                                    Just Extern
e  -> do
                                          let inst :: Instance
inst   = Annote -> Text -> Text -> [Natural] -> Instance
Instance Annote
an' Text
i Text
x [Natural]
cs
                                              drives :: [Stmt]
drives = [ Annote -> Text -> Text -> Exp -> Stmt
SInstIn Annote
an' Text
i Text
p Exp
arg | ((Text
p, Size
_), Exp
arg) <- [(Text, Size)] -> [Exp] -> [((Text, Size), Exp)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Extern -> [(Text, Size)]
extInputs Extern
e) [Exp]
es' ]
                                              result :: Exp
result = [Exp] -> Exp
cat [ Annote -> Size -> Text -> Exp
Var Annote
an' Size
psz (Text
i Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"." Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
q) | (Text
q, Size
psz) <- Extern -> [(Text, Size)]
extOutputs Extern
e ]
                                          (Exp, [Instance], [Stmt]) -> StateT Int m (Exp, [Instance], [Stmt])
forall a. a -> StateT Int m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
result, Instance
inst Instance -> [Instance] -> [Instance]
forall a. a -> [a] -> [a]
: [Instance]
i1, [Stmt] -> [Stmt]
forall a. [a] -> [a]
reverse [Stmt]
drives [Stmt] -> [Stmt] -> [Stmt]
forall a. Semigroup a => a -> a -> a
<> [Stmt]
s1)
                        | Bool
otherwise -> do
                              (es', i1, s1) <- Annote
-> HashMap Text Exp
-> ([Instance], [Stmt])
-> [Exp]
-> StateT Int m ([Exp], [Instance], [Stmt])
hoistExps Annote
an HashMap Text Exp
sub ([Instance], [Stmt])
acc [Exp]
es
                              pure (XCall an' sz x cs es', i1, s1)

            hoistExps :: Annote -> HashMap Name Exp -> ([Instance], [Stmt]) -> [Exp] -> StateT Int m ([Exp], [Instance], [Stmt])
            hoistExps :: Annote
-> HashMap Text Exp
-> ([Instance], [Stmt])
-> [Exp]
-> StateT Int m ([Exp], [Instance], [Stmt])
hoistExps Annote
an HashMap Text Exp
sub ([Instance], [Stmt])
acc = (([Exp], [Instance], [Stmt])
 -> Exp -> StateT Int m ([Exp], [Instance], [Stmt]))
-> ([Exp], [Instance], [Stmt])
-> [Exp]
-> StateT Int m ([Exp], [Instance], [Stmt])
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
                  (\ ([Exp]
es', [Instance]
i1, [Stmt]
s1) Exp
e -> (\ (Exp
e', [Instance]
i2, [Stmt]
s2) -> ([Exp]
es' [Exp] -> [Exp] -> [Exp]
forall a. Semigroup a => a -> a -> a
<> [Exp
e'], [Instance]
i2, [Stmt]
s2)) ((Exp, [Instance], [Stmt]) -> ([Exp], [Instance], [Stmt]))
-> StateT Int m (Exp, [Instance], [Stmt])
-> StateT Int m ([Exp], [Instance], [Stmt])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Annote
-> HashMap Text Exp
-> ([Instance], [Stmt])
-> Exp
-> StateT Int m (Exp, [Instance], [Stmt])
hoistExp Annote
an HashMap Text Exp
sub ([Instance]
i1, [Stmt]
s1) Exp
e)
                  ([], ([Instance], [Stmt]) -> [Instance]
forall a b. (a, b) -> a
fst ([Instance], [Stmt])
acc, ([Instance], [Stmt]) -> [Stmt]
forall a b. (a, b) -> b
snd ([Instance], [Stmt])
acc)