```
module DDC.Core.Transform.ANormal
(anormalise)
where
import DDC.Core.Exp
import qualified DDC.Type.Exp as T
import qualified DDC.Type.Compounds as T
import qualified DDC.Core.Transform.AnonymizeX as A
import qualified DDC.Core.Transform.LiftX as L

import qualified Data.Map as Map

-- **** Recording arities of known values
-- So we can try to create apps to fully apply

-- | Arities of known bound variables.
-- We need to track everything even if it's not a function to keep indices correct.
-- Just use zero for unknown/irrelevant
type Arities n = (Map.Map n Int, [Int])

-- | Empty arities context
arEmpty :: Ord n => Arities n
arEmpty = (Map.empty, [])

-- | Extend map with multiple bindings and their arities
arExtends :: Ord n => Arities n -> [(Bind n, Int)] -> Arities n
arExtends arity exts = foldl go arity exts
where	go (named,anon) (BNone _t,   _)    = (named,anon)
go (named,anon) (BAnon _t,   a)    = (named, a:anon)
go (named,anon) (BName n _t, a) = (Map.insert n a named, anon)

-- | Look up a binder's arity
arGet :: Ord n => Arities n -> Bound n -> Int
arGet (_named, anon) (UIx ix _)	  = anon !! ix
arGet (named, _anon) (UName n _)  = named Map.! n
-- Get a primitive's arity from its type
arGet (_named,_anon) (UPrim _ t)  = arityOfType t

-- **** Finding arities of expressions etc

-- | Count all the arrows and foralls, ignoring any effects
-- We can be sure that primitives don't effect until they're fully applied
arityOfType :: Ord n => Type n -> Int
arityOfType (T.TForall _ t)
=  1 + arityOfType t
arityOfType t
=  let (args, _) = T.takeTFunArgResult t in
length args

-- | Find arity of an expression. Count lambdas, use type for primitives
arityOfExp :: Ord n => Exp a n -> Int
-- Counting all binders, because they all correspond to XApps.
arityOfExp (XLam _ _ e)
= 1 + arityOfExp e
arityOfExp (XLAM _ _ e)
= 1 + arityOfExp e
-- Find primitive's constructor's arities from type,
-- we might need to do this for user defined constructors too.
arityOfExp (XCon _ (UPrim _ t))
= arityOfType t
-- Anything else we'll need to apply one at a time
arityOfExp _
= 0

-- | Retrieve binders from case pattern, so we can extend the arity context.
-- We don't know anything about their values, so record as 0.
aritiesOfPat :: Ord n => Pat n -> [(Bind n, Int)]
aritiesOfPat PDefault = []
aritiesOfPat (PData _b bs) = zip bs (repeat 0)

-- **** Actually converting to a-normal form

-- | Recursively transform expression into a-normal
anormal :: Ord n
=> Arities n	-- ^ environment, arities of bound variables
-> Exp a n	-- ^ expression to transform
-> [(Exp a n,a)]-- ^ arguments being applied to current expression
-> Exp a n

-- Application: just record argument and descend into function
anormal ar (XApp a lhs rhs) args
=  -- normalise rhs and add to arguments
let args' = (anormal ar rhs [], a) : args in
-- descend into lhs, remembering all args
anormal ar lhs args'

-- Anything other than application: if we're applied to arguments add bindings,
-- otherwise just recurse.
anormal ar x args
=  let x' = go x in
case args of
-- if there are no args, we're done
[] -> x'
-- there are arguments. we must apply them.
_  -> flattenLets \$ makeLets ar x' args
where
-- helper for descent
down ars e = anormal (arExtends ar ars) e []

-- we know x isn't an app.
go (XApp{}) = error "DDC.Core.Transform.ANormal.anormal: impossible XApp!"

-- leafy ones
go (XVar{}) = x
go (XCon{}) = x
go (XType{}) = x
go (XWitness{}) = x

-- lambdas
go (XLAM a b e) =
XLAM a b (down [(b,0)] e)
go (XLam a b e) =
XLam a b (down [(b,0)] e)

-- non-recursive let
go (XLet a (LLet m b le) re) =
let le' = down [] le in
let re' = down [(b, arityOfExp le')] re in
XLet a (LLet m b le') re'

-- recursive let
go (XLet a (LRec lets) re) =
let bs = map fst lets in
let es = map snd lets in
let ars= zip bs (map arityOfExp es) in
let es'= map (down ars) es in
let re'= down ars re in
XLet a (LRec \$ zip bs es') re'

-- letregion, just make sure we record bindings with dummy val
go (XLet a (LLetRegion b bs) re) =
let ars = zip bs (repeat 0) in
XLet a (LLetRegion b bs) (down ars re)

-- withregion: I don't think this should ever show up.
go (XLet a (LWithRegion b) re) =
XLet a (LWithRegion b) (down [] re)

-- case
go (XCase a e alts) =
let e' = down [] e in
let alts' = map (\(AAlt pat ae) -> AAlt pat (down (aritiesOfPat pat) ae)) alts in
XCase a e' alts'

-- cast
go (XCast a c e) =
XCast a c (down [] e)

-- | Convert an expression into a-normal form
anormalise :: Ord n => Exp a n -> Exp a n
anormalise x = anormal arEmpty x []

-- | Check if an expression needs a binding, or if it's simple enough to be applied as-is
isNormal :: Ord n => Exp a n -> Bool
-- Trivial expressions
isNormal (XVar{}) = True
isNormal (XCon{}) = True
isNormal (XType{}) = True
isNormal (XWitness{}) = True
-- Casts are ignored by code generator, so we can leave them in if their subexpression is normal
isNormal (XCast _ _ x) = isNormal x
isNormal _ = False

-- | Create lets for any non-trivial arguments
makeLets :: Ord n
=> Arities n	-- ^ environment, arities of bound variables
-> Exp a n	-- ^ function
-> [(Exp a n,a)]-- ^ arguments being applied to current expression
-> Exp a n
makeLets _  f0 [] = f0
makeLets ar f0 args@((_,annot):_) = go 0 (findArity f0) ((f0,annot):args) []
where
tBot = T.tBot T.kData

-- out of arguments, create XApps out of leftovers
go i _arf [] acc = mkApps i 0 acc
-- f is fully applied and we have arguments left to add:
--	create let for intermediate result
go i arf ((x,a):xs) acc | length acc > arf
=  XLet a (LLet LetStrict (BAnon tBot) (mkApps i 0 acc))
(go i 1 ((x,a):xs) [(XVar a \$ UIx 0 tBot,a)])
-- application to variable, don't bother binding
go i arf ((x,a):xs) acc | isNormal x
=  go i arf xs ((x,a):acc)
-- non-trivial argument, create binding
go i arf ((x,a):xs) acc
=  XLet a (LLet LetStrict (BAnon tBot) (L.liftX i x))
(go (i+1) arf xs ((x,a):acc))

-- fold list into applications
-- can't create empty app
mkApps _ _ []
= error "DDC.Core.Transform.ANormal.makeLets.mkApps: impossible empty list"

-- single element - this is the function
mkApps l _ [(x,_)] | isNormal x
= L.liftX l x
mkApps _ i [(_,a)]
= XVar a \$ UIx i tBot

-- apply this argument and recurse
mkApps l i ((x,a):xs) | isNormal x
= XApp a (mkApps l i xs) (L.liftX l x)
mkApps l i ((_,a):xs)
= XApp a (mkApps l (i+1) xs) (XVar a \$ UIx i tBot)

findArity (XVar _ b) = max (arGet ar b) 1
findArity x          = max (arityOfExp x) 1

-- | Perform let-floating on strict non-recursive lets
-- Only does the top level, to clean up the ones directly produced by makeLets.
-- let b1 = (let b2 = def2 in x2)
-- in x1
-- ==>
-- let b2 = def2
-- in let b1 = x2
-- in x1
flattenLets :: Ord n
=> Exp a n
-> Exp a n

-- We only do this if b2 is anonymous (ones generated by makeLets are).
-- If we tried to wrap x1 in b2 when b2's name is already used,
-- we'd be in trouble.
flattenLets
(XLet a1
(LLet LetStrict b1
(XLet a2 (LLet LetStrict b2@(BAnon _) def2) x2))
x1)
=  -- If b1 is anon, we don't want to lift references to it
let liftDepth = case b1 of { BAnon _ -> 1; _ -> 0 } in
let x1'	  = L.liftAtDepthX 1 liftDepth x1 in
XLet a2 (LLet LetStrict b2 def2) \$
flattenLets \$ XLet a1 (LLet LetStrict b1 x2) x1'

-- Same as above but b2 isn't anonymous - anonymize inner let & re-flatten.
flattenLets
(XLet a1
(LLet LetStrict b1 inner@(XLet _ (LLet LetStrict _ _) _))
x1)
=  flattenLets \$
XLet a1
(LLet LetStrict b1 (A.anonymizeX inner))
x1

-- Any let, its bound expression doesn't contain a strict non-recursive let so just flatten the body
flattenLets (XLet a1 llet1 x1)
=  XLet a1 llet1 (flattenLets x1)

-- Anything else we can ignore. We don't need to recurse, because this is always called immediately after makeLets.
flattenLets x = x
```