```
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.Type.Universe as U
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

-- I did have these as Maybe Int, but I think for our purposes 0==Nothing is fine
type Arities n = (Map.Map n Int, [Int])

arEmpty :: Ord n => Arities n
arEmpty = (Map.empty, [])

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)

arGet :: Ord n => Arities n -> Bound n -> Int
-- TODO unsafe ix
arGet (_named, anon) (UIx ix _)	  = anon !! ix
arGet (named, _anon) (UName n _)  = named Map.! n
-- Get a primitive's arity from its type.
-- Assuming all the primitives defer effects until fully applied.
arGet (_named,_anon) (UPrim _ t)  = arityOfType t

-- **** Finding arities of expressions etc

-- Count all the arrows, ignoring any effects
arityOfType :: Ord n => Type n -> Int
arityOfType (T.TForall _ t)
=  1 + arityOfType t
arityOfType t
=  let (args, _) = T.takeTFunArgResult t in
length args

arityOfExp :: Ord n => Exp a n -> Int
arityOfExp (XLam _ b e)
-- only count data binders
| isBinderData b
= 1 + arityOfExp e
arityOfExp (XLam _ _ e)
= 1 + arityOfExp e
arityOfExp (XLAM _ _ e)
= 1 + arityOfExp e
arityOfExp (XCon _ (UPrim _ t))
= arityOfType t
arityOfExp _
= 0

isBinderData :: Ord n => Bind n -> Bool
isBinderData b | Just U.UniverseData <- U.universeFromType1 (T.typeOfBind b)
=  True
isBinderData _ = False

-- We don't know anything about their values,
-- but we need to record them as 0 anyway (shadowing, de bruijn)
aritiesOfPat :: Ord n => Pat n -> [(Bind n, Int)]
aritiesOfPat PDefault = []
aritiesOfPat (PData _b bs) = zip bs (repeat 0)

-- **** Actually converting to a-normal form
anormal :: Ord n => Arities n -> Exp a n -> [Exp a n] -> Exp a n
anormal ar (XApp _ lhs rhs) args
=  -- normalise applicand and record arguments
let args' = anormal ar rhs [] : args in
-- descend into lhs, remembering all args
anormal ar lhs args'

anormal ar x args
=  let x' = go x in
-- if there are no args, we're done
case args of
[] -> x'
_  -> -- there are arguments. we must apply them.
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 "ANormal.anormal: impossible XApp!"

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

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)

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

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'

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

-- | (under development)
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 just be applied
isNormal :: Ord n => Exp a n -> Bool
isNormal (XVar{}) = True
isNormal (XCon{}) = True
isNormal (XType{}) = True
isNormal (XWitness{}) = True
isNormal (XCast _ _ x) = isNormal x
isNormal _ = False

makeLets ar f0 args = go 0 (findArity f0) (f0:args) []
where
tBot = T.tBot T.kData

-- sending arity of f to this is a hack because we should really be building up ar ctx?
go i _arf []  acc = mkApps i 0 acc
-- f is fully applied, and we *do* have arguments left to add
go i arf (x:xs) acc | length acc > arf
=  XLet (annotOf x) (LLet LetStrict (BAnon tBot) (mkApps i 0 acc))
(go i 1 (x:xs) [XVar (annotOf x) \$ UIx 0 tBot])
-- application to variable, don't bother binding
go i arf (x:xs) acc | isNormal x
=  go i arf xs (x:acc)
-- create binding
go i arf (x:xs) acc
=  XLet (annotOf x) (LLet LetStrict (BAnon tBot) (L.liftX i x))
(go (i+1) arf xs (x:acc))

mkApps _ _ []
= error "ANormal.makeLets.mkApps: impossible empty list"
mkApps l _ [x] | isNormal x
= L.liftX l x
mkApps _ i [x]
= XVar (annotOf x) \$ UIx i tBot

mkApps l i (x:xs) | isNormal x
= XApp (annotOf x) (mkApps l i xs) (L.liftX l x)
mkApps l i (x:xs)
= XApp (annotOf x) (mkApps l (i+1) xs) (XVar (annotOf x) \$ UIx i tBot)

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

-- does this exist elsewhere? ought it?
annotOf :: Exp a n -> a
annotOf (XVar a _) = a
annotOf (XCon a _) = a
annotOf (XApp a _ _) = a
annotOf (XLAM a _ _) = a
annotOf (XLam a _ _) = a
annotOf (XLet a _ _) = a
annotOf (XCase a _ _) = a
annotOf (XCast a _ _) = a
annotOf (XType{}) = error "DDC.Core.Transform.ANormal.annotOf: XType"
annotOf (XWitness{}) = error "DDC.Core.Transform.ANormal.annotOf: XWitness"

```