10 | module Compiler.LambdaLift

 12 | import Core.CompileExpr

 13 | import Core.Context

 15 | import Data.Vect

 17 | import Libraries.Data.SnocList.SizeOf

 18 | import Libraries.Data.List.Extra

 20 | %default covering

 24 |   ||| Type representing the syntax tree of an Idris 2 program after lambda

 32 |   public export

 33 |   data Lifted : Scoped where

 35 |        ||| A local variable in the lambda-lifted syntax tree.

 41 |        LLocal : {idx : Nat} -> FC -> (0 p : IsVar x idx vars) -> Lifted vars

 43 |        ||| A known function applied to exactly the right number of arguments.

 52 |        LAppName : FC -> (lazy : Maybe LazyReason) -> (n : Name) ->

 53 |                   (args : List (Lifted vars)) -> Lifted vars

 55 |        ||| A known function applied to fewer arguments than its arity.

 64 |        LUnderApp : FC -> (n : Name) -> (missing : Nat) ->

 65 |                    (args : List (Lifted vars)) -> Lifted vars

 67 |        ||| A closure applied to one more argument.

 77 |        LApp : FC -> (lazy : Maybe LazyReason) -> (closure : Lifted vars) ->

 78 |               (arg : Lifted vars) -> Lifted vars

 80 |        ||| A let binding: binding a new name to an existing expression.

 93 |        LLet : FC -> (x : Name) -> (expr : Lifted vars) ->

 94 |               (body : Lifted (x :: vars)) -> Lifted vars

 96 |        ||| Use of a constructor to construct a compound data type value.

103 |        LCon : FC -> (n : Name) -> (info : ConInfo) -> (tag : Maybe Int) ->

104 |               (args : List (Lifted vars)) -> Lifted vars

106 |        ||| An operator applied to operands.

113 |        LOp : {arity : _} ->

114 |              FC -> (lazy : Maybe LazyReason) -> (op : PrimFn arity) ->

115 |              (args : Vect arity (Lifted vars)) -> Lifted vars

117 |        ||| A second, more involved, form of primitive operation, defined using

127 |        LExtPrim : FC -> (lazy : Maybe LazyReason) -> (p : Name) ->

128 |                   (args : List (Lifted vars)) -> Lifted vars

130 |        ||| A case split on constructor tags.

136 |        LConCase : FC -> (expr : Lifted vars) ->

137 |                   (alts : List (LiftedConAlt vars)) ->

138 |                   (def : Maybe (Lifted vars)) -> Lifted vars

140 |        ||| A case split on a constant expression.

146 |        LConstCase : FC -> (expr : Lifted vars) ->

147 |                     (alts : List (LiftedConstAlt vars)) ->

148 |                     (def : Maybe (Lifted vars)) -> Lifted vars

150 |        ||| A primitive (constant) value.

151 |        LPrimVal : FC -> Constant -> Lifted vars

153 |        ||| An erased value.

154 |        LErased : FC -> Lifted vars

156 |        ||| A forceful crash of the program.

164 |        LCrash : FC -> (msg : String) -> Lifted vars

166 |   ||| A branch of an "LCon" (constructor tag) case statement.

170 |   public export

171 |   data LiftedConAlt : Scoped where

173 |        ||| Constructs a branch of an "LCon" (constructor tag) case statement.

188 |        MkLConAlt : (n : Name) -> (info : ConInfo) -> (tag : Maybe Int) ->

189 |                    (args : List Name) -> (body : Lifted (args ++ vars)) ->

190 |                    LiftedConAlt vars

192 |   ||| A branch of an "LConst" (constant expression) case statement.

196 |   public export

197 |   data LiftedConstAlt : Scoped where

199 |        ||| Constructs a branch of an "LConst" (constant expression) case

205 |        MkLConstAlt : (expr : Constant) -> (body : Lifted vars) ->

206 |                      LiftedConstAlt vars

210 | public export

211 | data LiftedDef : Type where

213 |      ||| Constructs a function definition.

219 |      -- We take the outer scope and the function arguments separately so that

220 |      -- we don't have to reshuffle de Bruijn indices, which is expensive.

221 |      -- This should be compiled as a function which takes 'args' first,

222 |      -- then 'reverse scope'.

223 |      -- (Sorry for the awkward API - it's to do with how the indices are

224 |      -- arranged for the variables, and it could be expensive to reshuffle them!

225 |      -- See Compiler.ANF for an example of how they get resolved to names)

226 |      MkLFun : (args : Scope) -> (scope : Scope) ->

227 |               (body : Lifted (Scope.addInner args scope)) -> LiftedDef

229 |      ||| Constructs a definition of a constructor for a compound data type.

238 |      MkLCon : (tag : Maybe Int) -> (arity : Nat) -> (nt : Maybe Nat) ->

239 |               LiftedDef

241 |      ||| Constructs a forward declaration of a foreign function.

248 |      MkLForeign : (ccs : List String) ->

249 |                   (fargs : List CFType) ->

250 |                   (ret : CFType) ->

251 |                   LiftedDef

253 |      ||| Constructs an error condition.

260 |      MkLError : (expl : Lifted Scope.empty) -> LiftedDef

262 | showLazy : Maybe LazyReason -> String

263 | showLazy = maybe "" $ (" " ++) . show

266 |   export

267 |   covering

268 |   {vs : _} -> Show (Lifted vs) where

269 |     show (LLocal {idx} _ p) = "!" ++ show (nameAt p)

270 |     show (LAppName fc lazy n args)

271 |         = show n ++ showLazy lazy ++ "(" ++ showSep ", " (map show args) ++ ")"

272 |     show (LUnderApp fc n m args)

273 |         = "<" ++ show n ++ " underapp " ++ show m ++ ">(" ++

274 |           showSep ", " (map show args) ++ ")"

275 |     show (LApp fc lazy c arg)

276 |         = show c ++ showLazy lazy ++ " @ (" ++ show arg ++ ")"

277 |     show (LLet fc x val sc)

278 |         = "%let " ++ show x ++ " = " ++ show val ++ " in " ++ show sc

279 |     show (LCon fc n _ t args)

280 |         = "%con " ++ show n ++ "(" ++ showSep ", " (map show args) ++ ")"

281 |     show (LOp fc lazy op args)

282 |         = "%op " ++ show op ++ showLazy lazy ++ "(" ++ showSep ", " (toList (map show args)) ++ ")"

283 |     show (LExtPrim fc lazy p args)

284 |         = "%extprim " ++ show p ++ showLazy lazy ++ "(" ++ showSep ", " (map show args) ++ ")"

285 |     show (LConCase fc sc alts def)

286 |         = "%case " ++ show sc ++ " of { "

287 |              ++ showSep "| " (map show alts) ++ " " ++ show def

288 |     show (LConstCase fc sc alts def)

289 |         = "%case " ++ show sc ++ " of { "

290 |              ++ showSep "| " (map show alts) ++ " " ++ show def

291 |     show (LPrimVal _ x) = show x

292 |     show (LErased _) = "___"

293 |     show (LCrash _ x) = "%CRASH(" ++ show x ++ ")"

295 |   export

296 |   covering

297 |   {vs : _} -> Show (LiftedConAlt vs) where

298 |     show (MkLConAlt n _ t args sc)

299 |         = "%conalt " ++ show n ++

300 |              "(" ++ showSep ", " (map show args) ++ ") => " ++ show sc

302 |   export

303 |   covering

304 |   {vs : _} -> Show (LiftedConstAlt vs) where

305 |     show (MkLConstAlt c sc)

306 |         = "%constalt(" ++ show c ++ ") => " ++ show sc

310 | Show LiftedDef where

311 |   show (MkLFun args scope exp)

312 |       = show args ++ show (reverse scope) ++ ": " ++ show exp

313 |   show (MkLCon tag arity pos)

314 |       = "Constructor tag " ++ show tag ++ " arity " ++ show arity ++

315 |         maybe "" (\n => " (newtype by " ++ show n ++ ")") pos

316 |   show (MkLForeign ccs args ret)

317 |       = "Foreign call " ++ show ccs ++ " " ++

318 |         show args ++ " -> " ++ show ret

319 |   show (MkLError exp) = "Error: " ++ show exp

322 | data Lifts : Type where

324 | record LDefs where

325 |   constructor MkLDefs

326 |   basename : Name -- top level name we're lifting from

327 |   defs : List (Name, LiftedDef) -- new definitions we made

328 |   nextName : Int -- name of next definition to lift

330 | genName : {auto l : Ref Lifts LDefs} ->

331 |           Core Name

333 |     = do ldefs <- get Lifts

334 |          let i = nextName ldefs

335 |          put Lifts ({ nextName := i + 1 } ldefs)

336 |          pure $ mkName (basename ldefs) i

337 |   where

338 |     mkName : Name -> Int -> Name

339 |     mkName (NS ns b) i = NS ns (mkName b i)

340 |     mkName (UN n) i = MN (displayUserName n) i

341 |     mkName (DN _ n) i = mkName n i

342 |     mkName (CaseBlock outer inner) i = MN ("case block in " ++ outer ++ " (" ++ show inner ++ ")") i

343 |     mkName (WithBlock outer inner) i = MN ("with block in " ++ outer ++ " (" ++ show inner ++ ")") i

344 |     mkName n i = MN (show n) i

346 | unload : FC -> (lazy : Maybe LazyReason) -> Lifted vars -> List (Lifted vars) -> Core (Lifted vars)

347 | unload fc _ f [] = pure f

349 | unload fc lazy f (a :: as) = unload fc Nothing (LApp fc lazy f a) as

351 | record Used (vars : Scope) where

352 |   constructor MkUsed

353 |   used : Vect (length vars) Bool

355 | initUsed : {vars : _} -> Used vars

356 | initUsed {vars} = MkUsed (replicate (length vars) False)

358 | weakenUsed : {outer : _} -> Used vars -> Used (outer ++ vars)

359 | weakenUsed {outer} (MkUsed xs) =

360 |   MkUsed (rewrite lengthDistributesOverAppend outer vars in

361 |          (replicate (length outer) False ++ xs))

363 | contractUsed : (Used (x::vars)) -> Used vars

364 | contractUsed (MkUsed xs) = MkUsed (tail xs)

366 | contractUsedMany : {remove : _} ->

367 |                    (Used (remove ++ vars)) ->

368 |                    Used vars

369 | contractUsedMany {remove=[]} x = x

370 | contractUsedMany {remove=(r::rs)} x = contractUsedMany {remove=rs} (contractUsed x)

372 | markUsed : {vars : _} ->

373 |            (idx : Nat) ->

374 |            {0 prf : IsVar x idx vars} ->

375 |            Used vars ->

376 |            Used vars

377 | markUsed {vars} {prf} idx (MkUsed us) =

378 |   let newUsed = replaceAt (finIdx prf) True us in

379 |   MkUsed newUsed

380 |     where

381 |     finIdx : {vars : _} -> {idx : _} ->

382 |              (0 prf : IsVar x idx vars) ->

383 |              Fin (length vars)

384 |     finIdx {idx=Z} First = FZ

385 |     finIdx {idx=S x} (Later l) = FS (finIdx l)

387 | getUnused : Used vars ->

388 |             Vect (length vars) Bool

389 | getUnused (MkUsed uv) = map not uv

392 | dropped : (vars : Scope) ->

393 |           (drop : Vect (length vars) Bool) ->

394 |           Scope

395 | dropped [] _ = []

396 | dropped (x::xs) (False::us) = x::(dropped xs us)

397 | dropped (x::xs) (True::us) = dropped xs us

399 | usedVars : {vars : _} ->

400 |            {auto l : Ref Lifts LDefs} ->

401 |            Used vars ->

402 |            Lifted vars ->

403 |            Used vars

404 | usedVars used (LLocal {idx} fc prf) =

405 |   markUsed {prf} idx used

406 | usedVars used (LAppName fc lazy n args) =

407 |   foldl (usedVars {vars}) used args

408 | usedVars used (LUnderApp fc n miss args) =

409 |   foldl (usedVars {vars}) used args

410 | usedVars used (LApp fc lazy c arg) =

411 |   usedVars (usedVars used arg) c

412 | usedVars used (LLet fc x val sc) =

413 |   let innerUsed = contractUsed $ usedVars (weakenUsed {outer=Scope.single x} used) sc in

414 |       usedVars innerUsed val

415 | usedVars used (LCon fc n ci tag args) =

416 |   foldl (usedVars {vars}) used args

417 | usedVars used (LOp fc lazy fn args) =

418 |   foldl (usedVars {vars}) used args

419 | usedVars used (LExtPrim fc lazy fn args) =

420 |   foldl (usedVars {vars}) used args

421 | usedVars used (LConCase fc sc alts def) =

422 |     let defUsed = maybe used (usedVars used {vars}) def

423 |         scDefUsed = usedVars defUsed sc in

424 |         foldl usedConAlt scDefUsed alts

425 |   where

426 |     usedConAlt : {default Nothing lazy : Maybe LazyReason} ->

427 |                   Used vars -> LiftedConAlt vars -> Used vars

428 |     usedConAlt used (MkLConAlt n ci tag args sc) =

429 |       contractUsedMany {remove=args} (usedVars (weakenUsed used) sc)

431 | usedVars used (LConstCase fc sc alts def) =

432 |     let defUsed = maybe used (usedVars used {vars}) def

433 |         scDefUsed = usedVars defUsed sc in

434 |         foldl usedConstAlt scDefUsed alts

435 |   where

436 |     usedConstAlt : {default Nothing lazy : Maybe LazyReason} ->

437 |                     Used vars -> LiftedConstAlt vars -> Used vars

438 |     usedConstAlt used (MkLConstAlt c sc) = usedVars used sc

439 | usedVars used (LPrimVal {}) = used

440 | usedVars used (LErased {})  = used

441 | usedVars used (LCrash {})   = used

443 | dropIdx : {vars : _} ->

444 |           {idx : _} ->

445 |           (outer : Scope) ->

446 |           (unused : Vect (length vars) Bool) ->

447 |           (0 p : IsVar x idx (outer ++ vars)) ->

448 |           Var (outer ++ (dropped vars unused))

449 | dropIdx [] (False::_) First = first

450 | dropIdx [] (True::_) First = assert_total $

451 |   idris_crash "INTERNAL ERROR: Referenced variable marked as unused"

452 | dropIdx [] (False::rest) (Later p) = Var.later $ dropIdx Scope.empty rest p

453 | dropIdx [] (True::rest) (Later p) = dropIdx Scope.empty rest p

454 | dropIdx (_::xs) unused First = first

455 | dropIdx (_::xs) unused (Later p) = Var.later $ dropIdx xs unused p

457 | dropUnused : {vars : _} ->

458 |              {auto _ : Ref Lifts LDefs} ->

459 |              {outer : Scope} ->

460 |              (unused : Vect (length vars) Bool) ->

461 |              (l : Lifted (outer ++ vars)) ->

462 |              Lifted (outer ++ (dropped vars unused))

463 | dropUnused _ (LPrimVal fc val) = LPrimVal fc val

464 | dropUnused _ (LErased fc) = LErased fc

465 | dropUnused _ (LCrash fc msg) = LCrash fc msg

466 | dropUnused {outer} unused (LLocal fc p) =

467 |   let (MkVar p') = dropIdx outer unused p in LLocal fc p'

468 | dropUnused unused (LCon fc n ci tag args) =

469 |   let args' = map (dropUnused unused) args in

470 |       LCon fc n ci tag args'

471 | dropUnused {outer} unused (LLet fc n val sc) =

472 |   let val' = dropUnused unused val

473 |       sc' = dropUnused {outer=n::outer} (unused) sc in

474 |       LLet fc n val' sc'

475 | dropUnused unused (LApp fc lazy c arg) =

476 |   let c' = dropUnused unused c

477 |       arg' = dropUnused unused arg in

478 |       LApp fc lazy c' arg'

479 | dropUnused unused (LOp fc lazy fn args) =

480 |   let args' = map (dropUnused unused) args in

481 |       LOp fc lazy fn args'

482 | dropUnused unused (LExtPrim fc lazy n args) =

483 |   let args' = map (dropUnused unused) args in

484 |       LExtPrim fc lazy n args'

485 | dropUnused unused (LAppName fc lazy n args) =

486 |   let args' = map (dropUnused unused) args in

487 |       LAppName fc lazy n args'

488 | dropUnused unused (LUnderApp fc n miss args) =

489 |   let args' = map (dropUnused unused) args in

490 |       LUnderApp fc n miss args'

491 | dropUnused {vars} {outer} unused (LConCase fc sc alts def) =

492 |   let alts' = map dropConCase alts in

493 |       LConCase fc (dropUnused unused sc) alts' (map (dropUnused unused) def)

494 |   where

495 |     dropConCase : LiftedConAlt (outer ++ vars) ->

496 |                   LiftedConAlt (outer ++ (dropped vars unused))

497 |     dropConCase (MkLConAlt n ci t args sc) =

498 |       let sc' = (rewrite sym $ appendAssociative args outer vars in sc)

499 |           droppedSc = dropUnused {vars=vars} {outer=args++outer} unused sc' in

500 |       MkLConAlt n ci t args (rewrite appendAssociative args outer (dropped vars unused) in droppedSc)

501 | dropUnused {vars} {outer} unused (LConstCase fc sc alts def) =

502 |   let alts' = map dropConstCase alts in

503 |       LConstCase fc (dropUnused unused sc) alts' (map (dropUnused unused) def)

504 |   where

505 |     dropConstCase : LiftedConstAlt (outer ++ vars) ->

506 |                     LiftedConstAlt (outer ++ (dropped vars unused))

507 |     dropConstCase (MkLConstAlt c val) = MkLConstAlt c (dropUnused unused val)

510 |   makeLam : {vars : _} ->

511 |             {auto l : Ref Lifts LDefs} ->

512 |             {doLazyAnnots : Bool} ->

513 |             {default Nothing lazy : Maybe LazyReason} ->

514 |             FC -> (bound : Scope) ->

515 |             CExp (bound ++ vars) -> Core (Lifted vars)

516 |   makeLam fc bound (CLam _ x sc') = makeLam fc {doLazyAnnots} {lazy} (x :: bound) sc'

517 |   makeLam {vars} fc bound sc

518 |       = do scl <- liftExp {doLazyAnnots} {lazy} sc

519 |            -- Find out which variables aren't used in the new definition, and

520 |            -- do not abstract over them in the new definition.

521 |            let scUsedL = usedVars initUsed scl

522 |                unusedContracted = contractUsedMany {remove=bound} scUsedL

523 |                unused = getUnused unusedContracted

524 |                scl' = dropUnused {outer=bound} unused scl

525 |            n <- genName

526 |            update Lifts { defs $= ((n, MkLFun (dropped vars unused) bound scl') ::) }

527 |            pure $ LUnderApp fc n (length bound) (allVars fc vars unused)

528 |     where

530 |         allPrfs : (vs : Scope) -> SizeOf seen ->

531 |                   (unused : Vect (length vs) Bool) ->

532 |                   List (Var (seen <>> vs))

533 |         allPrfs [] _ _ = []

534 |         allPrfs (v :: vs) p (False::uvs) = mkVarChiply p :: allPrfs vs (p :< _) uvs

535 |         allPrfs (v :: vs) p (True::uvs) = allPrfs vs (p :< _) uvs

537 |         -- apply to all the variables. 'First' will be first in the last, which

538 |         -- is good, because the most recently bound name is the first argument to

539 |         -- the resulting function

540 |         allVars : FC -> (vs : Scope) -> (unused : Vect (length vs) Bool) -> List (Lifted vs)

541 |         allVars fc vs unused = map (\ (MkVar p) => LLocal fc p) (allPrfs vs [<] unused)

545 |   liftExp : {vars : _} ->

546 |             {auto l : Ref Lifts LDefs} ->

547 |             {doLazyAnnots : Bool} ->

548 |             {default Nothing lazy : Maybe LazyReason} ->

549 |             CExp vars -> Core (Lifted vars)

550 |   liftExp (CLocal fc prf) = pure $ LLocal fc prf

551 |   liftExp (CRef fc n) = pure $ LAppName fc lazy n [] -- probably shouldn't happen!

552 |   liftExp (CLam fc x sc) = makeLam {doLazyAnnots} {lazy} fc (Scope.single x) sc

553 |   liftExp (CLet fc x _ val sc) = pure $ LLet fc x !(liftExp {doLazyAnnots} val) !(liftExp {doLazyAnnots} sc)

554 |   liftExp (CApp fc (CRef _ n) args) -- names are applied exactly in compileExp

555 |       = pure $ LAppName fc lazy n !(traverse (liftExp {doLazyAnnots}) args)

556 |   liftExp (CApp fc f args)

557 |       = unload fc lazy !(liftExp {doLazyAnnots} f) !(traverse (liftExp {doLazyAnnots}) args)

558 |   liftExp (CCon fc n ci t args) = pure $ LCon fc n ci t !(traverse (liftExp {doLazyAnnots}) args)

559 |   liftExp (COp fc op args)

560 |       = pure $ LOp fc lazy op !(traverseArgs args)

561 |     where

562 |       traverseArgs : Vect n (CExp vars) -> Core (Vect n (Lifted vars))

563 |       traverseArgs [] = pure []

564 |       traverseArgs (a :: as) = pure $ !(liftExp {doLazyAnnots} a) :: !(traverseArgs as)

565 |   liftExp (CExtPrim fc p args) = pure $ LExtPrim fc lazy p !(traverse (liftExp {doLazyAnnots}) args)

566 |   liftExp (CForce fc lazy tm) = if doLazyAnnots

567 |     then liftExp {doLazyAnnots} {lazy = Nothing} tm

568 |     else liftExp {doLazyAnnots} (CApp fc tm [CErased fc])

569 |   liftExp (CDelay fc lazy tm) = if doLazyAnnots

570 |     then liftExp {doLazyAnnots} {lazy = Just lazy} tm

571 |     else liftExp {doLazyAnnots} (CLam fc (MN "act" 0) (weaken tm))

572 |   liftExp (CConCase fc sc alts def)

573 |       = pure $ LConCase fc !(liftExp {doLazyAnnots} sc) !(traverse (liftConAlt {lazy}) alts)

574 |                            !(traverseOpt (liftExp {doLazyAnnots}) def)

575 |     where

576 |       liftConAlt : {default Nothing lazy : Maybe LazyReason} ->

577 |                    CConAlt vars -> Core (LiftedConAlt vars)

578 |       liftConAlt (MkConAlt n ci t args sc) = pure $ MkLConAlt n ci t args !(liftExp {doLazyAnnots} {lazy} sc)

579 |   liftExp (CConstCase fc sc alts def)

580 |       = pure $ LConstCase fc !(liftExp {doLazyAnnots} sc) !(traverse liftConstAlt alts)

581 |                              !(traverseOpt (liftExp {doLazyAnnots}) def)

582 |     where

583 |       liftConstAlt : {default Nothing lazy : Maybe LazyReason} ->

584 |                      CConstAlt vars -> Core (LiftedConstAlt vars)

585 |       liftConstAlt (MkConstAlt c sc) = pure $ MkLConstAlt c !(liftExp {doLazyAnnots} {lazy} sc)

586 |   liftExp (CPrimVal fc c) = pure $ LPrimVal fc c

587 |   liftExp (CErased fc) = pure $ LErased fc

588 |   liftExp (CCrash fc str) = pure $ LCrash fc str

591 | liftBody : {vars : _} -> {doLazyAnnots : Bool} ->

592 |            Name -> CExp vars -> Core (Lifted vars, List (Name, LiftedDef))

593 | liftBody n tm

594 |     = do l <- newRef Lifts (MkLDefs n [] 0)

595 |          tml <- liftExp {doLazyAnnots} {l} tm

596 |          ldata <- get Lifts

597 |          pure (tml, defs ldata)

600 | lambdaLiftDef : (doLazyAnnots : Bool) -> Name -> CDef -> Core (List (Name, LiftedDef))

601 | lambdaLiftDef doLazyAnnots n (MkFun args exp)

602 |     = do (expl, defs) <- liftBody {doLazyAnnots} n exp

603 |          pure ((n, MkLFun args Scope.empty expl) :: defs)

604 | lambdaLiftDef _ n (MkCon t a nt) = pure [(n, MkLCon t a nt)]

605 | lambdaLiftDef _ n (MkForeign ccs fargs ty) = pure [(n, MkLForeign ccs fargs ty)]

606 | lambdaLiftDef doLazyAnnots n (MkError exp)

607 |     = do (expl, defs) <- liftBody {doLazyAnnots} n exp

608 |          pure ((n, MkLError expl) :: defs)

616 | lambdaLift :  (doLazyAnnots : Bool)

617 |            -> (Name,FC,CDef)

618 |            -> Core (List (Name, LiftedDef))

619 | lambdaLift doLazyAnnots (n,_,def) = lambdaLiftDef doLazyAnnots n def