module Lift where

import Data.Functor.Foldable
import Data.List
import Types

converge f a = let a' = f a in if a' == a then a else converge f a'

factorial :: AExp
factorial =
    ( Lam
        ["n"]
        ( Let
            "m"
            (Call "factorial" [(ASub (Ident "n") (Number 1))])
            ( If
                (ALt (Ident "n") (Number 2))
                (FC (Atom (Number 1)))
                (FC (Atom (AMul (Ident "m") (Ident "n"))))
            )
        )
    )

liftArgs :: AExp -> AExp
liftArgs lam@(Lam args body) = Lam (args ++ cata findVarsAExp lam) body
liftArgs rest = rest

-- | F-Algebras to find free variables
findVars :: CExpF [String] -> [String]
findVars (LetF ident fc rest) = converge (\\ [ident]) $ rest ++ (findVarsFC fc)
findVars (IfF cond t e) = t ++ e ++ (cata findVarsAExp cond)
findVars (FCF fc) = findVarsFC fc

findVarsFC :: Funcall -> [String]
findVarsFC (Atom aexp) = cata findVarsAExp aexp
findVarsFC (Call id args) = id : (args >>= cata findVarsAExp)

findVarsAExp :: AExpF [String] -> [String]
findVarsAExp (IdentF ns) = [ns]
findVarsAExp (LamF args cexp) = converge (\\ args) (fold findVars cexp)
findVarsAExp def = foldMap id def