module Compiler.ANF import Compiler.LambdaLift import Core.CompileExprimport Core.Contextimport Core.Coreimport Core.TT import Data.Vect %default covering -- Convert the lambda lifted form to ANF, with variable names made explicit.-- i.e. turn intermediate expressions into let bindings. Every argument is-- a variable as a result. mutual  public export  data AVar : Type where       ALocal : Int -> AVar       ANull : AVar   public export  data ANF : Type where    AV : FC -> AVar -> ANF    AAppName : FC -> Name -> List AVar -> ANF    AUnderApp : FC -> Name -> (missing : Nat) -> (args : List AVar) -> ANF    AApp : FC -> (closure : AVar) -> (arg : AVar) -> ANF    ALet : FC -> (var : Int) -> ANF -> ANF -> ANF    ACon : FC -> Name -> (tag : Maybe Int) -> List AVar -> ANF    AOp : FC -> PrimFn arity -> Vect arity AVar -> ANF    AExtPrim : FC -> Name -> List AVar -> ANF    AConCase : FC -> AVar -> List AConAlt -> Maybe ANF -> ANF    AConstCase : FC -> AVar -> List AConstAlt -> Maybe ANF -> ANF    APrimVal : FC -> Constant -> ANF    AErased : FC -> ANF    ACrash : FC -> String -> ANF   public export  data AConAlt : Type where       MkAConAlt : Name -> (tag : Maybe Int) -> (args : List Int) ->                   ANF -> AConAlt   public export  data AConstAlt : Type where       MkAConstAlt : Constant -> ANF -> AConstAlt public exportdata ANFDef : Type where     MkAFun : (args : List Int) -> ANF -> ANFDef     MkACon : (tag : Maybe Int) -> (arity : Nat) -> ANFDef     MkAForeign : (ccs : List String) -> (fargs : List CFType) ->                  CFType -> ANFDef     MkAError : ANF -> ANFDef mutual  export  Show AVar where    show (ALocal i) = "v" ++ show i    show ANull = "[__]"   export  Show ANF where    show (AV _ v) = show v    show (AAppName fc n args)        = show n ++ "(" ++ showSep ", " (map show args) ++ ")"    show (AUnderApp fc n m args)        = "<" ++ show n ++ " underapp " ++ show m ++ ">(" ++           showSep ", " (map show args) ++ ")"    show (AApp fc c arg)        = show c ++ " @ (" ++ show arg ++ ")"    show (ALet fc x val sc)        = "%let v" ++ show x ++ " = " ++ show val ++ " in " ++ show sc    show (ACon fc n t args)        = "%con " ++ show n ++ "(" ++ showSep ", " (map show args) ++ ")"    show (AOp fc op args)        = "%op " ++ show op ++ "(" ++ showSep ", " (toList (map show args)) ++ ")"    show (AExtPrim fc p args)        = "%extprim " ++ show p ++ "(" ++ showSep ", " (map show args) ++ ")"    show (AConCase fc sc alts def)        = "%case " ++ show sc ++ " of { "             ++ showSep "| " (map show alts) ++ " " ++ show def    show (AConstCase fc sc alts def)        = "%case " ++ show sc ++ " of { "             ++ showSep "| " (map show alts) ++ " " ++ show def    show (APrimVal _ x) = show x    show (AErased _) = "___"    show (ACrash _ x) = "%CRASH(" ++ show x ++ ")"   export  Show AConAlt where    show (MkAConAlt n t args sc)        = "%conalt " ++ show n ++              "(" ++ showSep ", " (map showArg args) ++ ") => " ++ show sc      where        showArg : Int -> String        showArg i = "v" ++ show i   export  Show AConstAlt where    show (MkAConstAlt c sc)        = "%constalt(" ++ show c ++ ") => " ++ show sc exportShow ANFDef where  show (MkAFun args exp) = show args ++ ": " ++ show exp  show (MkACon tag arity)      = "Constructor tag " ++ show tag ++ " arity " ++ show arity  show (MkAForeign ccs args ret)      = "Foreign call " ++ show ccs ++ " " ++        show args ++ " -> " ++ show ret  show (MkAError exp) = "Error: " ++ show exp data AVars : List Name -> Type where     Nil : AVars []     (::) : Int -> AVars xs -> AVars (x :: xs) data Next : Type where nextVar : {auto v : Ref Next Int} ->          Core IntnextVar    = do i <- get Next         put Next (i + 1)         pure i lookup : IsVar x idx vs -> AVars vs -> Intlookup First (x :: xs) = xlookup (Later p) (x :: xs) = lookup p xs bindArgs : {auto v : Ref Next Int} ->           List ANF -> Core (List (AVar, Maybe ANF))bindArgs [] = pure []bindArgs (AV fc var :: xs)    = do xs' <- bindArgs xs         pure $ (var, Nothing) :: xs'bindArgs (AErased fc :: xs)    = do xs' <- bindArgs xs         pure $ (ANull, Nothing) :: xs'bindArgs (x :: xs)    = do i <- nextVar         xs' <- bindArgs xs         pure $ (ALocal i, Just x) :: xs' letBind : {auto v : Ref Next Int} ->          FC -> List ANF -> (List AVar -> ANF) -> Core ANFletBind fc args f    = do bargs <- bindArgs args         pure $ doBind [] bargs  where    doBind : List AVar -> List (AVar, Maybe ANF) -> ANF    doBind vs [] = f (reverse vs)    doBind vs ((ALocal i, Just t) :: xs)        = ALet fc i t (doBind (ALocal i :: vs) xs)    doBind vs ((var, _) :: xs) = doBind (var :: vs) xs toVect : (n : Nat) -> List a -> Maybe (Vect n a)toVect Z [] = Just []toVect (S k) (x :: xs)    = do xs' <- toVect k xs         pure (x :: xs')toVect _ _ = Nothing mlet : {auto v : Ref Next Int} ->       FC -> ANF -> (AVar -> ANF) -> Core ANFmlet fc (AV _ var) sc = pure $ sc varmlet fc val sc    = do i <- nextVar         pure $ ALet fc i val (sc (ALocal i))  mutual  anfArgs : {auto v : Ref Next Int} ->            FC -> AVars vars ->            List (Lifted vars) -> (List AVar -> ANF) -> Core ANF  anfArgs fc vs args f      = do args' <- traverse (anf vs) args           letBind fc args' f   anf : {auto v : Ref Next Int} ->        AVars vars -> Lifted vars -> Core ANF  anf vs (LLocal fc p) = pure $ AV fc (ALocal (lookup p vs))  anf vs (LAppName fc n args)      = anfArgs fc vs args (AAppName fc n)  anf vs (LUnderApp fc n m args)      = anfArgs fc vs args (AUnderApp fc n m)  anf vs (LApp fc f a)      = anfArgs fc vs [f, a]                (\args => case args of                               [fvar, avar] => AApp fc fvar avar                               _ => ACrash fc "Can't happen (AApp)")  anf vs (LLet fc x val sc)      = do i <- nextVar           let vs' = i :: vs           pure $ ALet fc i !(anf vs val) !(anf vs' sc)  anf vs (LCon fc n t args)      = anfArgs fc vs args (ACon fc n t)  anf vs (LOp {arity} fc op args)      = do args' <- traverse (anf vs) (toList args)           letBind fc args'                (\args => case toVect arity args of                               Nothing => ACrash fc "Can't happen (AOp)"                               Just argsv => AOp fc op argsv)  anf vs (LExtPrim fc p args)      = anfArgs fc vs args (AExtPrim fc p)  anf vs (LConCase fc scr alts def)      = do scr' <- anf vs scr           alts' <- traverse (anfConAlt vs) alts           def' <- traverseOpt (anf vs) def           mlet fc scr' (\x => AConCase fc x alts' def')  anf vs (LConstCase fc scr alts def)      = do scr' <- anf vs scr           alts' <- traverse (anfConstAlt vs) alts           def' <- traverseOpt (anf vs) def           mlet fc scr' (\x => AConstCase fc x alts' def')  anf vs (LPrimVal fc c) = pure $ APrimVal fc c  anf vs (LErased fc) = pure $ AErased fc  anf vs (LCrash fc err) = pure $ ACrash fc err   anfConAlt : {auto v : Ref Next Int} ->              AVars vars -> LiftedConAlt vars -> Core AConAlt  anfConAlt vs (MkLConAlt n t args sc)      = do (is, vs') <- bindArgs args vs           pure $ MkAConAlt n t is !(anf vs' sc)    where      bindArgs : {auto v : Ref Next Int} ->                 (args : List Name) -> AVars vars' ->                 Core (List Int, AVars (args ++ vars'))      bindArgs [] vs = pure ([], vs)      bindArgs (n :: ns) vs          = do i <- nextVar               (is, vs') <- bindArgs ns vs               pure (i :: is, i :: vs')   anfConstAlt : {auto v : Ref Next Int} ->                AVars vars -> LiftedConstAlt vars -> Core AConstAlt  anfConstAlt vs (MkLConstAlt c sc)      = pure $ MkAConstAlt c !(anf vs sc) exporttoANF : LiftedDef -> Core ANFDeftoANF (MkLFun args scope sc)    = do v <- newRef Next (the Int 0)         (iargs, vsNil) <- bindArgs args []         let vs : AVars args = rewrite sym (appendNilRightNeutral args) in                                      vsNil         (iargs', vs) <- bindArgs scope vs         pure $ MkAFun (iargs ++ reverse iargs') !(anf vs sc)  where    bindArgs : {auto v : Ref Next Int} ->               (args : List Name) -> AVars vars' ->               Core (List Int, AVars (args ++ vars'))    bindArgs [] vs = pure ([], vs)    bindArgs (n :: ns) vs        = do i <- nextVar             (is, vs') <- bindArgs ns vs             pure (i :: is, i :: vs')toANF (MkLCon t a ns) = pure $ MkACon t atoANF (MkLForeign ccs fargs t) = pure $ MkAForeign ccs fargs ttoANF (MkLError err)    = do v <- newRef Next (the Int 0)         pure $ MkAError !(anf [] err)