10 | module Compiler.LambdaLift
12 | import Core.CompileExpr
17 | import Libraries.Data.SnocList.SizeOf
32 | data Lifted : Scoped where
40 | LLocal : {idx : Nat} -> FC -> (0 p : IsVar x idx vars) -> Lifted vars
51 | LAppName : FC -> (lazy : Maybe LazyReason) -> (n : Name) ->
52 | (args : List (Lifted vars)) -> Lifted vars
63 | LUnderApp : FC -> (n : Name) -> (missing : Nat) ->
64 | (args : List (Lifted vars)) -> Lifted vars
76 | LApp : FC -> (lazy : Maybe LazyReason) -> (closure : Lifted vars) ->
77 | (arg : Lifted vars) -> Lifted vars
92 | LLet : FC -> (x : Name) -> (expr : Lifted vars) ->
93 | (body : Lifted (x :: vars)) -> Lifted vars
102 | LCon : FC -> (n : Name) -> (info : ConInfo) -> (tag : Maybe Int) ->
103 | (args : List (Lifted vars)) -> Lifted vars
112 | LOp : {arity : _} ->
113 | FC -> (lazy : Maybe LazyReason) -> (op : PrimFn arity) ->
114 | (args : Vect arity (Lifted vars)) -> Lifted vars
126 | LExtPrim : FC -> (lazy : Maybe LazyReason) -> (p : Name) ->
127 | (args : List (Lifted vars)) -> Lifted vars
135 | LConCase : FC -> (expr : Lifted vars) ->
136 | (alts : List (LiftedConAlt vars)) ->
137 | (def : Maybe (Lifted vars)) -> Lifted vars
145 | LConstCase : FC -> (expr : Lifted vars) ->
146 | (alts : List (LiftedConstAlt vars)) ->
147 | (def : Maybe (Lifted vars)) -> Lifted vars
150 | LPrimVal : FC -> Constant -> Lifted vars
153 | LErased : FC -> Lifted vars
163 | LCrash : FC -> (msg : String) -> Lifted vars
170 | data LiftedConAlt : Scoped where
187 | MkLConAlt : (n : Name) -> (info : ConInfo) -> (tag : Maybe Int) ->
188 | (args : List Name) -> (body : Lifted (args ++ vars)) ->
196 | data LiftedConstAlt : Scoped where
204 | MkLConstAlt : (expr : Constant) -> (body : Lifted vars) ->
205 | LiftedConstAlt vars
210 | data LiftedDef : Type where
225 | MkLFun : (args : Scope) -> (scope : Scope) ->
226 | (body : Lifted (Scope.addInner args scope)) -> LiftedDef
237 | MkLCon : (tag : Maybe Int) -> (arity : Nat) -> (nt : Maybe Nat) ->
247 | MkLForeign : (ccs : List String) ->
248 | (fargs : List CFType) ->
259 | MkLError : (expl : Lifted Scope.empty) -> LiftedDef
261 | showLazy : Maybe LazyReason -> String
262 | showLazy = maybe "" $
(" " ++) . show
267 | {vs : _} -> Show (Lifted vs) where
268 | show (LLocal {idx} _ p) = "!" ++ show (nameAt p)
269 | show (LAppName fc lazy n args)
270 | = show n ++ showLazy lazy ++ "(" ++ showSep ", " (map show args) ++ ")"
271 | show (LUnderApp fc n m args)
272 | = "<" ++ show n ++ " underapp " ++ show m ++ ">(" ++
273 | showSep ", " (map show args) ++ ")"
274 | show (LApp fc lazy c arg)
275 | = show c ++ showLazy lazy ++ " @ (" ++ show arg ++ ")"
276 | show (LLet fc x val sc)
277 | = "%let " ++ show x ++ " = " ++ show val ++ " in " ++ show sc
278 | show (LCon fc n _ t args)
279 | = "%con " ++ show n ++ "(" ++ showSep ", " (map show args) ++ ")"
280 | show (LOp fc lazy op args)
281 | = "%op " ++ show op ++ showLazy lazy ++ "(" ++ showSep ", " (toList (map show args)) ++ ")"
282 | show (LExtPrim fc lazy p args)
283 | = "%extprim " ++ show p ++ showLazy lazy ++ "(" ++ showSep ", " (map show args) ++ ")"
284 | show (LConCase fc sc alts def)
285 | = "%case " ++ show sc ++ " of { "
286 | ++ showSep "| " (map show alts) ++ " " ++ show def
287 | show (LConstCase fc sc alts def)
288 | = "%case " ++ show sc ++ " of { "
289 | ++ showSep "| " (map show alts) ++ " " ++ show def
290 | show (LPrimVal _ x) = show x
291 | show (LErased _) = "___"
292 | show (LCrash _ x) = "%CRASH(" ++ show x ++ ")"
296 | {vs : _} -> Show (LiftedConAlt vs) where
297 | show (MkLConAlt n _ t args sc)
298 | = "%conalt " ++ show n ++
299 | "(" ++ showSep ", " (map show args) ++ ") => " ++ show sc
303 | {vs : _} -> Show (LiftedConstAlt vs) where
304 | show (MkLConstAlt c sc)
305 | = "%constalt(" ++ show c ++ ") => " ++ show sc
309 | Show LiftedDef where
310 | show (MkLFun args scope exp)
311 | = show args ++ show (reverse scope) ++ ": " ++ show exp
312 | show (MkLCon tag arity pos)
313 | = "Constructor tag " ++ show tag ++ " arity " ++ show arity ++
314 | maybe "" (\n => " (newtype by " ++ show n ++ ")") pos
315 | show (MkLForeign ccs args ret)
316 | = "Foreign call " ++ show ccs ++ " " ++
317 | show args ++ " -> " ++ show ret
318 | show (MkLError exp) = "Error: " ++ show exp
321 | data Lifts : Type where
324 | constructor MkLDefs
326 | defs : List (Name, LiftedDef)
329 | genName : {auto l : Ref Lifts LDefs} ->
332 | = do ldefs <- get Lifts
333 | let i = nextName ldefs
334 | put Lifts ({ nextName := i + 1 } ldefs)
335 | pure $
mkName (basename ldefs) i
337 | mkName : Name -> Int -> Name
338 | mkName (NS ns b) i = NS ns (mkName b i)
339 | mkName (UN n) i = MN (displayUserName n) i
340 | mkName (DN _ n) i = mkName n i
341 | mkName (CaseBlock outer inner) i = MN ("case block in " ++ outer ++ " (" ++ show inner ++ ")") i
342 | mkName (WithBlock outer inner) i = MN ("with block in " ++ outer ++ " (" ++ show inner ++ ")") i
343 | mkName n i = MN (show n) i
345 | unload : FC -> (lazy : Maybe LazyReason) -> Lifted vars -> List (Lifted vars) -> Core (Lifted vars)
346 | unload fc _ f [] = pure f
348 | unload fc lazy f (a :: as) = unload fc Nothing (LApp fc lazy f a) as
350 | record Used (vars : Scope) where
352 | used : Vect (length vars) Bool
354 | initUsed : {vars : _} -> Used vars
355 | initUsed {vars} = MkUsed (replicate (length vars) False)
358 | lengthDistributesOverAppend
359 | : (xs, ys : List a)
360 | -> length (xs ++ ys) = length xs + length ys
361 | lengthDistributesOverAppend [] ys = Refl
362 | lengthDistributesOverAppend (x :: xs) ys =
363 | cong S $
lengthDistributesOverAppend xs ys
365 | weakenUsed : {outer : _} -> Used vars -> Used (outer ++ vars)
366 | weakenUsed {outer} (MkUsed xs) =
367 | MkUsed (rewrite lengthDistributesOverAppend outer vars in
368 | (replicate (length outer) False ++ xs))
370 | contractUsed : (Used (x::vars)) -> Used vars
371 | contractUsed (MkUsed xs) = MkUsed (tail xs)
373 | contractUsedMany : {remove : _} ->
374 | (Used (remove ++ vars)) ->
376 | contractUsedMany {remove=[]} x = x
377 | contractUsedMany {remove=(r::rs)} x = contractUsedMany {remove=rs} (contractUsed x)
379 | markUsed : {vars : _} ->
381 | {0 prf : IsVar x idx vars} ->
384 | markUsed {vars} {prf} idx (MkUsed us) =
385 | let newUsed = replaceAt (finIdx prf) True us in
388 | finIdx : {vars : _} -> {idx : _} ->
389 | (0 prf : IsVar x idx vars) ->
391 | finIdx {idx=Z} First = FZ
392 | finIdx {idx=S x} (Later l) = FS (finIdx l)
394 | getUnused : Used vars ->
395 | Vect (length vars) Bool
396 | getUnused (MkUsed uv) = map not uv
399 | dropped : (vars : Scope) ->
400 | (drop : Vect (length vars) Bool) ->
403 | dropped (x::xs) (False::us) = x::(dropped xs us)
404 | dropped (x::xs) (True::us) = dropped xs us
406 | usedVars : {vars : _} ->
407 | {auto l : Ref Lifts LDefs} ->
411 | usedVars used (LLocal {idx} fc prf) =
412 | markUsed {prf} idx used
413 | usedVars used (LAppName fc lazy n args) =
414 | foldl (usedVars {vars}) used args
415 | usedVars used (LUnderApp fc n miss args) =
416 | foldl (usedVars {vars}) used args
417 | usedVars used (LApp fc lazy c arg) =
418 | usedVars (usedVars used arg) c
419 | usedVars used (LLet fc x val sc) =
420 | let innerUsed = contractUsed $
usedVars (weakenUsed {outer=Scope.single x} used) sc in
421 | usedVars innerUsed val
422 | usedVars used (LCon fc n ci tag args) =
423 | foldl (usedVars {vars}) used args
424 | usedVars used (LOp fc lazy fn args) =
425 | foldl (usedVars {vars}) used args
426 | usedVars used (LExtPrim fc lazy fn args) =
427 | foldl (usedVars {vars}) used args
428 | usedVars used (LConCase fc sc alts def) =
429 | let defUsed = maybe used (usedVars used {vars}) def
430 | scDefUsed = usedVars defUsed sc in
431 | foldl usedConAlt scDefUsed alts
433 | usedConAlt : {default Nothing lazy : Maybe LazyReason} ->
434 | Used vars -> LiftedConAlt vars -> Used vars
435 | usedConAlt used (MkLConAlt n ci tag args sc) =
436 | contractUsedMany {remove=args} (usedVars (weakenUsed used) sc)
438 | usedVars used (LConstCase fc sc alts def) =
439 | let defUsed = maybe used (usedVars used {vars}) def
440 | scDefUsed = usedVars defUsed sc in
441 | foldl usedConstAlt scDefUsed alts
443 | usedConstAlt : {default Nothing lazy : Maybe LazyReason} ->
444 | Used vars -> LiftedConstAlt vars -> Used vars
445 | usedConstAlt used (MkLConstAlt c sc) = usedVars used sc
446 | usedVars used (LPrimVal {}) = used
447 | usedVars used (LErased {}) = used
448 | usedVars used (LCrash {}) = used
450 | dropIdx : {vars : _} ->
453 | (unused : Vect (length vars) Bool) ->
454 | (0 p : IsVar x idx (outer ++ vars)) ->
455 | Var (outer ++ (dropped vars unused))
456 | dropIdx [] (False::_) First = first
457 | dropIdx [] (True::_) First = assert_total $
458 | idris_crash "INTERNAL ERROR: Referenced variable marked as unused"
459 | dropIdx [] (False::rest) (Later p) = Var.later $
dropIdx Scope.empty rest p
460 | dropIdx [] (True::rest) (Later p) = dropIdx Scope.empty rest p
461 | dropIdx (_::xs) unused First = first
462 | dropIdx (_::xs) unused (Later p) = Var.later $
dropIdx xs unused p
464 | dropUnused : {vars : _} ->
465 | {auto _ : Ref Lifts LDefs} ->
467 | (unused : Vect (length vars) Bool) ->
468 | (l : Lifted (outer ++ vars)) ->
469 | Lifted (outer ++ (dropped vars unused))
470 | dropUnused _ (LPrimVal fc val) = LPrimVal fc val
471 | dropUnused _ (LErased fc) = LErased fc
472 | dropUnused _ (LCrash fc msg) = LCrash fc msg
473 | dropUnused {outer} unused (LLocal fc p) =
474 | let (MkVar p') = dropIdx outer unused p in LLocal fc p'
475 | dropUnused unused (LCon fc n ci tag args) =
476 | let args' = map (dropUnused unused) args in
477 | LCon fc n ci tag args'
478 | dropUnused {outer} unused (LLet fc n val sc) =
479 | let val' = dropUnused unused val
480 | sc' = dropUnused {outer=n::outer} (unused) sc in
482 | dropUnused unused (LApp fc lazy c arg) =
483 | let c' = dropUnused unused c
484 | arg' = dropUnused unused arg in
485 | LApp fc lazy c' arg'
486 | dropUnused unused (LOp fc lazy fn args) =
487 | let args' = map (dropUnused unused) args in
488 | LOp fc lazy fn args'
489 | dropUnused unused (LExtPrim fc lazy n args) =
490 | let args' = map (dropUnused unused) args in
491 | LExtPrim fc lazy n args'
492 | dropUnused unused (LAppName fc lazy n args) =
493 | let args' = map (dropUnused unused) args in
494 | LAppName fc lazy n args'
495 | dropUnused unused (LUnderApp fc n miss args) =
496 | let args' = map (dropUnused unused) args in
497 | LUnderApp fc n miss args'
498 | dropUnused {vars} {outer} unused (LConCase fc sc alts def) =
499 | let alts' = map dropConCase alts in
500 | LConCase fc (dropUnused unused sc) alts' (map (dropUnused unused) def)
502 | dropConCase : LiftedConAlt (outer ++ vars) ->
503 | LiftedConAlt (outer ++ (dropped vars unused))
504 | dropConCase (MkLConAlt n ci t args sc) =
505 | let sc' = (rewrite sym $
appendAssociative args outer vars in sc)
506 | droppedSc = dropUnused {vars=vars} {outer=args++outer} unused sc' in
507 | MkLConAlt n ci t args (rewrite appendAssociative args outer (dropped vars unused) in droppedSc)
508 | dropUnused {vars} {outer} unused (LConstCase fc sc alts def) =
509 | let alts' = map dropConstCase alts in
510 | LConstCase fc (dropUnused unused sc) alts' (map (dropUnused unused) def)
512 | dropConstCase : LiftedConstAlt (outer ++ vars) ->
513 | LiftedConstAlt (outer ++ (dropped vars unused))
514 | dropConstCase (MkLConstAlt c val) = MkLConstAlt c (dropUnused unused val)
517 | makeLam : {vars : _} ->
518 | {auto l : Ref Lifts LDefs} ->
519 | {doLazyAnnots : Bool} ->
520 | {default Nothing lazy : Maybe LazyReason} ->
521 | FC -> (bound : Scope) ->
522 | CExp (bound ++ vars) -> Core (Lifted vars)
523 | makeLam fc bound (CLam _ x sc') = makeLam fc {doLazyAnnots} {lazy} (x :: bound) sc'
524 | makeLam {vars} fc bound sc
525 | = do scl <- liftExp {doLazyAnnots} {lazy} sc
528 | let scUsedL = usedVars initUsed scl
529 | unusedContracted = contractUsedMany {remove=bound} scUsedL
530 | unused = getUnused unusedContracted
531 | scl' = dropUnused {outer=bound} unused scl
533 | update Lifts { defs $= ((n, MkLFun (dropped vars unused) bound scl') ::) }
534 | pure $
LUnderApp fc n (length bound) (allVars fc vars unused)
537 | allPrfs : (vs : Scope) -> SizeOf seen ->
538 | (unused : Vect (length vs) Bool) ->
539 | List (Var (seen <>> vs))
540 | allPrfs [] _ _ = []
541 | allPrfs (v :: vs) p (False::uvs) = mkVarChiply p :: allPrfs vs (p :< _) uvs
542 | allPrfs (v :: vs) p (True::uvs) = allPrfs vs (p :< _) uvs
547 | allVars : FC -> (vs : Scope) -> (unused : Vect (length vs) Bool) -> List (Lifted vs)
548 | allVars fc vs unused = map (\ (MkVar p) => LLocal fc p) (allPrfs vs [<] unused)
552 | liftExp : {vars : _} ->
553 | {auto l : Ref Lifts LDefs} ->
554 | {doLazyAnnots : Bool} ->
555 | {default Nothing lazy : Maybe LazyReason} ->
556 | CExp vars -> Core (Lifted vars)
557 | liftExp (CLocal fc prf) = pure $
LLocal fc prf
558 | liftExp (CRef fc n) = pure $
LAppName fc lazy n []
559 | liftExp (CLam fc x sc) = makeLam {doLazyAnnots} {lazy} fc (Scope.single x) sc
560 | liftExp (CLet fc x _ val sc) = pure $
LLet fc x !(liftExp {doLazyAnnots} val) !(liftExp {doLazyAnnots} sc)
561 | liftExp (CApp fc (CRef _ n) args)
562 | = pure $
LAppName fc lazy n !(traverse (liftExp {doLazyAnnots}) args)
563 | liftExp (CApp fc f args)
564 | = unload fc lazy !(liftExp {doLazyAnnots} f) !(traverse (liftExp {doLazyAnnots}) args)
565 | liftExp (CCon fc n ci t args) = pure $
LCon fc n ci t !(traverse (liftExp {doLazyAnnots}) args)
566 | liftExp (COp fc op args)
567 | = pure $
LOp fc lazy op !(traverseArgs args)
569 | traverseArgs : Vect n (CExp vars) -> Core (Vect n (Lifted vars))
570 | traverseArgs [] = pure []
571 | traverseArgs (a :: as) = pure $
!(liftExp {doLazyAnnots} a) :: !(traverseArgs as)
572 | liftExp (CExtPrim fc p args) = pure $
LExtPrim fc lazy p !(traverse (liftExp {doLazyAnnots}) args)
573 | liftExp (CForce fc lazy tm) = if doLazyAnnots
574 | then liftExp {doLazyAnnots} {lazy = Nothing} tm
575 | else liftExp {doLazyAnnots} (CApp fc tm [CErased fc])
576 | liftExp (CDelay fc lazy tm) = if doLazyAnnots
577 | then liftExp {doLazyAnnots} {lazy = Just lazy} tm
578 | else liftExp {doLazyAnnots} (CLam fc (MN "act" 0) (weaken tm))
579 | liftExp (CConCase fc sc alts def)
580 | = pure $
LConCase fc !(liftExp {doLazyAnnots} sc) !(traverse (liftConAlt {lazy}) alts)
581 | !(traverseOpt (liftExp {doLazyAnnots}) def)
583 | liftConAlt : {default Nothing lazy : Maybe LazyReason} ->
584 | CConAlt vars -> Core (LiftedConAlt vars)
585 | liftConAlt (MkConAlt n ci t args sc) = pure $
MkLConAlt n ci t args !(liftExp {doLazyAnnots} {lazy} sc)
586 | liftExp (CConstCase fc sc alts def)
587 | = pure $
LConstCase fc !(liftExp {doLazyAnnots} sc) !(traverse liftConstAlt alts)
588 | !(traverseOpt (liftExp {doLazyAnnots}) def)
590 | liftConstAlt : {default Nothing lazy : Maybe LazyReason} ->
591 | CConstAlt vars -> Core (LiftedConstAlt vars)
592 | liftConstAlt (MkConstAlt c sc) = pure $
MkLConstAlt c !(liftExp {doLazyAnnots} {lazy} sc)
593 | liftExp (CPrimVal fc c) = pure $
LPrimVal fc c
594 | liftExp (CErased fc) = pure $
LErased fc
595 | liftExp (CCrash fc str) = pure $
LCrash fc str
598 | liftBody : {vars : _} -> {doLazyAnnots : Bool} ->
599 | Name -> CExp vars -> Core (Lifted vars, List (Name, LiftedDef))
601 | = do l <- newRef Lifts (MkLDefs n [] 0)
602 | tml <- liftExp {doLazyAnnots} {l} tm
604 | pure (tml, defs ldata)
607 | lambdaLiftDef : (doLazyAnnots : Bool) -> Name -> CDef -> Core (List (Name, LiftedDef))
608 | lambdaLiftDef doLazyAnnots n (MkFun args exp)
609 | = do (expl, defs) <- liftBody {doLazyAnnots} n exp
610 | pure ((n, MkLFun args Scope.empty expl) :: defs)
611 | lambdaLiftDef _ n (MkCon t a nt) = pure [(n, MkLCon t a nt)]
612 | lambdaLiftDef _ n (MkForeign ccs fargs ty) = pure [(n, MkLForeign ccs fargs ty)]
613 | lambdaLiftDef doLazyAnnots n (MkError exp)
614 | = do (expl, defs) <- liftBody {doLazyAnnots} n exp
615 | pure ((n, MkLError expl) :: defs)
623 | lambdaLift : (doLazyAnnots : Bool)
625 | -> Core (List (Name, LiftedDef))
626 | lambdaLift doLazyAnnots (n,_,def) = lambdaLiftDef doLazyAnnots n def