0 | module TTImp.PartialEval

  2 | import Core.Env

  3 | import Core.Hash

  4 | import Core.Metadata

  5 | import Core.Value

  6 | import Core.UnifyState

  8 | import Idris.REPL.Opts

  9 | import Idris.Syntax

 11 | import TTImp.Elab.Check

 12 | import TTImp.TTImp

 13 | import TTImp.TTImp.Functor

 14 | import TTImp.TTImp.Traversals

 15 | import TTImp.Unelab

 17 | import Protocol.Hex

 19 | import Libraries.Data.NameMap

 20 | import Libraries.Data.NatSet

 21 | import Libraries.Data.WithDefault

 22 | import Libraries.Data.SnocList.SizeOf

 24 | %default covering

 26 | data ArgMode' tm = Static tm | Dynamic

 28 | ArgMode : Type

 29 | ArgMode = ArgMode' ClosedTerm

 31 | traverseArgMode : (a -> Core b) -> ArgMode' a -> Core (ArgMode' b)

 32 | traverseArgMode f (Static t) = Static <$> f t

 33 | traverseArgMode f Dynamic = pure Dynamic

 36 | Show a => Show (ArgMode' a) where

 37 |   show (Static tm) = "Static " ++ show tm

 38 |   show Dynamic = "Dynamic"

 41 | getStatic : ArgMode -> Maybe ClosedTerm

 42 | getStatic Dynamic = Nothing

 43 | getStatic (Static t) = Just t

 45 | specialiseTy : {vars : _} ->

 46 |                Nat -> List (Nat, ClosedTerm) -> Term vars -> Term vars

 47 | specialiseTy i specs (Bind fc x (Pi fc' c p ty) sc)

 48 |     = case lookup i specs of

 49 |            Nothing => Bind fc x (Pi fc' c p ty) $ -- easier later if everything explicit

 50 |                         specialiseTy (1 + i) specs sc

 51 |            Just tm => specialiseTy (1 + i) specs (subst (embed tm) sc)

 52 | specialiseTy i specs tm = tm

 54 | substLoc : {vs : _} ->

 55 |            Nat -> Term vs -> Term vs -> Term vs

 56 | substLoc i tm (Local fc l idx p)

 57 |     = if i == idx then tm else (Local fc l idx p)

 58 | substLoc i tm (Bind fc x b sc)

 59 |     = Bind fc x (map (substLoc i tm) b) (substLoc (1 + i) (weaken tm) sc)

 60 | substLoc i tm (Meta fc n r args)

 61 |     = Meta fc n r (map (substLoc i tm) args)

 62 | substLoc i tm (App fc f a) = App fc (substLoc i tm f) (substLoc i tm a)

 63 | substLoc i tm (As fc s f a) = As fc s (substLoc i tm f) (substLoc i tm a)

 64 | substLoc i tm (TDelayed fc r d) = TDelayed fc r (substLoc i tm d)

 65 | substLoc i tm (TDelay fc r ty d) = TDelay fc r (substLoc i tm ty) (substLoc i tm d)

 66 | substLoc i tm (TForce fc r d) = TForce fc r (substLoc i tm d)

 67 | substLoc i tm x = x

 69 | substLocs : {vs : _} ->

 70 |             List (Nat, Term vs) -> Term vs -> Term vs

 71 | substLocs [] tm = tm

 72 | substLocs ((i, tm') :: subs) tm = substLocs subs (substLoc i tm' tm)

 74 | mkSubsts : Nat -> List (Nat, ClosedTerm) ->

 75 |            List (Term vs) -> Term vs -> Maybe (List (Nat, Term vs))

 76 | mkSubsts i specs [] rhs = Just []

 77 | mkSubsts i specs (arg :: args) rhs

 78 |     = do subs <- mkSubsts (1 + i) specs args rhs

 79 |          case lookup i specs of

 80 |               Nothing => Just subs

 81 |               Just tm => case arg of

 82 |                               Local _ _ idx _ =>

 83 |                                    Just ((idx, embed tm) :: subs)

 84 |                               As _ _ (Local _ _ idx1 _) (Local _ _ idx2 _) =>

 85 |                                    Just ((idx1, embed tm) :: (idx2, embed tm) :: subs)

 86 |                               As _ _ _ (Local _ _ idx _) =>

 87 |                                    Just ((idx, embed tm) :: subs)

 88 |                               _ => Nothing

 92 | specPatByVar : List (Nat, ClosedTerm) ->

 93 |                 (vs ** (Env Term vs, Term vs, Term vs)) ->

 94 |                 Maybe (vs ** (Env Term vs, Term vs, Term vs))

 95 | specPatByVar specs (vs ** (env, lhs, rhs))

 96 |     = do let (fn, args) = getFnArgs lhs

 97 |          psubs <- mkSubsts 0 specs args rhs

 98 |          let lhs' = apply (getLoc fn) fn args

 99 |          pure (vs ** (env, substLocs psubs lhs', substLocs psubs rhs))

101 | specByVar : List (Nat, ClosedTerm) ->

102 |             List (vs ** (Env Term vs, Term vs, Term vs)) ->

103 |             Maybe (List (vs ** (Env Term vs, Term vs, Term vs)))

104 | specByVar specs [] = pure []

105 | specByVar specs (p :: ps)

106 |     = do p' <- specPatByVar specs p

107 |          ps' <- specByVar specs ps

108 |          pure (p' :: ps')

110 | dropSpec : Nat -> List (Nat, ClosedTerm) -> List a -> List a

111 | dropSpec i sargs [] = []

112 | dropSpec i sargs (x :: xs)

113 |     = case lookup i sargs of

114 |            Nothing => x :: dropSpec (1 + i) sargs xs

115 |            Just _ => dropSpec (1 + i) sargs xs

117 | getSpecPats : {auto c : Ref Ctxt Defs} ->

118 |               {auto u : Ref UST UState} ->

119 |               FC -> Name ->

120 |               (fn : Name) -> (stk : List (FC, Term vars)) ->

121 |               ClosedNF -> -- Type of 'fn'

122 |               List (Nat, ArgMode) -> -- All the arguments

123 |               List (Nat, ClosedTerm) -> -- Just the static ones

124 |               List (vs ** (Env Term vs, Term vs, Term vs)) ->

125 |               Core (Maybe (List ImpClause))

126 | getSpecPats fc pename fn stk fnty args sargs pats

127 |    = do -- First, see if all the specialised positions are variables. If so,

128 |         -- substitute the arguments directly into the RHS

129 |         let Nothing = specByVar sargs pats

130 |             | Just specpats =>

131 |                    do ps <- traverse (unelabPat pename) specpats

132 |                       pure (Just ps)

133 |         -- Otherwise, build a new definition by taking the remaining arguments

134 |         -- on the lhs, and using the specialised function application on the rhs.

135 |         -- Then, this will get evaluated on elaboration.

136 |         dynnames <- mkDynNames args

137 |         let lhs = apply (IVar fc pename) (map (IBindVar fc) dynnames)

138 |         rhs <- mkRHSargs fnty (IVar fc fn) dynnames args

139 |         pure (Just [PatClause fc lhs rhs])

140 |   where

141 |     mkDynNames : List (Nat, ArgMode) -> Core (List Name)

142 |     mkDynNames [] = pure []

143 |     mkDynNames ((_, Dynamic) :: as) = [| genVarName "_pe" :: mkDynNames as |]

144 |     mkDynNames (_ :: as) = mkDynNames as

146 |     -- Build a RHS from the type of the function to be specialised, the

147 |     -- dynamic argument names, and the list of given arguments. We assume

148 |     -- the latter two correspond appropriately.

149 |     mkRHSargs : ClosedNF -> RawImp -> List Name -> List (Nat, ArgMode) ->

150 |                 Core RawImp

151 |     mkRHSargs (NBind _ x (Pi _ _ Explicit _) sc) app (a :: as) ((_, Dynamic) :: ds)

152 |         = do defs <- get Ctxt

153 |              sc' <- sc defs (toClosure defaultOpts Env.empty (Erased fc Placeholder))

154 |              mkRHSargs sc' (IApp fc app (IVar fc a)) as ds

155 |     mkRHSargs (NBind _ x (Pi {}) sc) app (a :: as) ((_, Dynamic) :: ds)

156 |         = do defs <- get Ctxt

157 |              sc' <- sc defs (toClosure defaultOpts Env.empty (Erased fc Placeholder))

158 |              mkRHSargs sc' (INamedApp fc app x (IVar fc a)) as ds

159 |     mkRHSargs (NBind _ x (Pi _ _ Explicit _) sc) app as ((_, Static tm) :: ds)

160 |         = do defs <- get Ctxt

161 |              sc' <- sc defs (toClosure defaultOpts Env.empty (Erased fc Placeholder))

162 |              tm' <- unelabNoSugar Env.empty tm

163 |              mkRHSargs sc' (IApp fc app (map rawName tm')) as ds

164 |     mkRHSargs (NBind _ x (Pi _ _ Implicit _) sc) app as ((_, Static tm) :: ds)

165 |         = do defs <- get Ctxt

166 |              sc' <- sc defs (toClosure defaultOpts Env.empty (Erased fc Placeholder))

167 |              tm' <- unelabNoSugar Env.empty tm

168 |              mkRHSargs sc' (INamedApp fc app x (map rawName tm')) as ds

169 |     mkRHSargs (NBind _ _ (Pi _ _ AutoImplicit _) sc) app as ((_, Static tm) :: ds)

170 |         = do defs <- get Ctxt

171 |              sc' <- sc defs (toClosure defaultOpts Env.empty (Erased fc Placeholder))

172 |              tm' <- unelabNoSugar Env.empty tm

173 |              mkRHSargs sc' (IAutoApp fc app (map rawName tm')) as ds

174 |     -- Type will depend on the value here (we assume a variadic function) but

175 |     -- the argument names are still needed

176 |     mkRHSargs ty app (a :: as) ((_, Dynamic) :: ds)

177 |         = mkRHSargs ty (IApp fc app (IVar fc a)) as ds

178 |     mkRHSargs _ app _ _

179 |         = pure app

181 |     getRawArgs : List (Arg' Name) -> RawImp -> List (Arg' Name)

182 |     getRawArgs args (IApp fc f arg) = getRawArgs (Explicit fc arg :: args) f

183 |     getRawArgs args (INamedApp fc f n arg)

184 |         = getRawArgs (Named fc n arg :: args) f

185 |     getRawArgs args (IAutoApp fc f arg)

186 |         = getRawArgs (Auto fc arg :: args) f

187 |     getRawArgs args tm = args

189 |     reapply : RawImp -> List (Arg' Name) -> RawImp

190 |     reapply f [] = f

191 |     reapply f (Explicit fc arg :: args) = reapply (IApp fc f arg) args

192 |     reapply f (Named fc n arg :: args)

193 |         = reapply (INamedApp fc f n arg) args

194 |     reapply f (Auto fc arg :: args)

195 |         = reapply (IAutoApp fc f arg) args

197 |     dropArgs : Name -> RawImp -> RawImp

198 |     dropArgs pename tm = reapply (IVar fc pename) (dropSpec 0 sargs (getRawArgs [] tm))

200 |     unelabPat : Name -> (vs ** (Env Term vs, Term vs, Term vs)) ->

201 |                 Core ImpClause

202 |     unelabPat pename (_ ** (env, lhs, rhs))

203 |         = do logTerm "specialise" 20 "Unelaborating LHS:" lhs

204 |              lhsapp <- unelabNoSugar env lhs

205 |              log "specialise" 20 $ "Unelaborating LHS to: \{show lhsapp}"

206 |              let lhs' = dropArgs pename (map rawName lhsapp)

207 |              defs <- get Ctxt

208 |              rhs <- normaliseArgHoles defs env rhs

209 |              rhs <- unelabNoSugar env rhs

210 |              let rhs = flip mapTTImp rhs $ \case

211 |                       IHole fc _ => Implicit fc False

212 |                       tm => tm

213 |              pure (PatClause fc lhs' (map rawName rhs))

215 |     unelabLHS : Name -> (vs ** (Env Term vs, Term vs, Term vs)) ->

216 |                 Core RawImp

217 |     unelabLHS pename (_ ** (env, lhs, rhs))

218 |         = do lhsapp <- unelabNoSugar env lhs

219 |              pure $ dropArgs pename (map rawName lhsapp)

224 | getReducible : List Name -> -- calls to check

225 |                NameMap Nat -> -- which nodes have been visited. If the entry is

226 |                               -- present, it's visited

227 |                Defs -> Core (NameMap Nat)

228 | getReducible [] refs defs = pure refs

229 | getReducible (n :: rest) refs defs

230 |   = do let Nothing = lookup n refs

231 |            | Just _ => getReducible rest refs defs

232 |        case !(lookupCtxtExact n (gamma defs)) of

233 |             Nothing => getReducible rest refs defs

234 |             Just def =>

235 |               do let refs' = insert n 65536 refs

236 |                  let calls = refersTo def

237 |                  getReducible (keys calls ++ rest) refs' defs

239 | mkSpecDef : {auto c : Ref Ctxt Defs} ->

240 |             {auto m : Ref MD Metadata} ->

241 |             {auto u : Ref UST UState} ->

242 |             {auto s : Ref Syn SyntaxInfo} ->

243 |             {auto o : Ref ROpts REPLOpts} ->

244 |             FC -> GlobalDef ->

245 |             Name -> List (Nat, ArgMode) -> Name -> List (FC, Term vars) ->

246 |             Core (Term vars)

247 | mkSpecDef {vars} fc gdef pename sargs fn stk

248 |     = handleUnify {unResolve = True}

249 |        (do defs <- get Ctxt

250 |            setAllPublic True

251 |            let staticargs

252 |                  = mapMaybe (\ (x, s) => case s of

253 |                                               Dynamic => Nothing

254 |                                               Static t => Just (x, t)) sargs

255 |            let peapp = applyStackWithFC (Ref fc Func pename) (dropSpec 0 staticargs stk)

256 |            Nothing <- lookupCtxtExact pename (gamma defs)

257 |                | Just _ => -- already specialised

258 |                            do log "specialise" 5 $ "Already specialised " ++ show pename

259 |                               pure peapp

260 |            logC "specialise.declare" 5 $

261 |                    do fnfull <- toFullNames fn

262 |                       args' <- traverse (\ (i, arg) =>

263 |                                    do arg' <- the (Core ArgMode) $ case arg of

264 |                                                    Static a =>

265 |                                                       pure $ Static !(toFullNames a)

266 |                                                    Dynamic => pure Dynamic

267 |                                       pure (show (i, arg'))) sargs

268 |                       pure $ "Specialising " ++ show fnfull ++

269 |                              " (" ++ show fn ++ ") -> \{show pename} by " ++

270 |                              showSep ", " args'

271 |            let sty = specialiseTy 0 staticargs (type gdef)

272 |            logTermNF "specialise" 3 ("Specialised type " ++ show pename) Env.empty sty

274 |            -- Add as RigW - if it's something else, we don't need it at

275 |            -- runtime anyway so this is wasted effort, therefore a failure

276 |            -- is okay!

277 |            let defflags := propagateFlags (flags gdef)

278 |            log "specialise.flags" 20 "Defining \{show pename} with flags: \{show defflags}"

279 |            peidx <- addDef pename

280 |                   $ the (GlobalDef -> GlobalDef) { flags := defflags }

281 |                   $ newDef fc pename top Scope.empty sty (specified Public) None

282 |            addToSave (Resolved peidx)

284 |            -- Reduce the function to be specialised, and reduce any name in

285 |            -- the arguments at most once (so that recursive definitions aren't

286 |            -- unfolded forever)

287 |            let specnames = getAllRefs empty (map snd sargs)

288 |            specLimits <- traverse (\n => pure (n, 1))

289 |                                   (keys specnames)

291 |            defs <- get Ctxt

292 |            reds <- getReducible [fn] empty defs

293 |            setFlag fc (Resolved peidx) (PartialEval (specLimits ++ toList reds))

295 |            let PMDef pminfo pmargs ct tr pats = definition gdef

296 |                | _ => pure (applyStackWithFC (Ref fc Func fn) stk)

297 |            logC "specialise" 5 $

298 |                    do inpats <- traverse unelabDef pats

299 |                       pure $ "Attempting to specialise:\n" ++

300 |                              showSep "\n" (map showPat inpats)

302 |            Just newpats <- getSpecPats fc pename fn stk !(nf defs Env.empty (type gdef))

303 |                                        sargs staticargs pats

304 |                 | Nothing => pure (applyStackWithFC (Ref fc Func fn) stk)

305 |            log "specialise" 5 $ "New patterns for " ++ show pename ++ ":\n" ++

306 |                     showSep "\n" (map showPat newpats)

307 |            processDecl [InPartialEval] (MkNested []) Env.empty

308 |                        (IDef fc (Resolved peidx) newpats)

309 |            setAllPublic False

310 |            pure peapp)

311 |            -- If the partially evaluated definition fails, just use the initial

312 |            -- application. It might indicates a bug in the P.E. function generation

313 |            -- if it fails, but I don't want the whole system to be dependent on

314 |            -- the correctness of PE!

315 |         (\err =>

316 |            do logC "specialise.fail" 1 $ do

317 |                  fn <- toFullNames fn

318 |                  pure "Partial evaluation of \{show fn} failed:\n\{show err}"

319 |               update Ctxt { peFailures $= insert pename () }

320 |               pure (applyStackWithFC (Ref fc Func fn) stk))

321 |   where

323 |     identityFlag : List (Nat, ArgMode) -> Nat -> Maybe Nat

324 |     identityFlag [] k = pure k

325 |     identityFlag ((pos, mode) :: sargs) k

326 |       = k <$ guard (k < pos)

327 |       <|> (case mode of { Static _ => (`minus` 1); Dynamic => id }) <$> identityFlag sargs k

330 |     propagateFlags : List DefFlag -> List DefFlag

331 |     propagateFlags = mapMaybe $ \case

332 |       Deprecate => Nothing

333 |       Overloadable => Nothing

334 |       Identity k => Identity <$> identityFlag sargs k

335 |       fl => Just fl

337 |     getAllRefs : NameMap Bool -> List ArgMode -> NameMap Bool

338 |     getAllRefs ns (Dynamic :: xs) = getAllRefs ns xs

339 |     getAllRefs ns (Static t :: xs)

340 |         = addRefs False (UN Underscore) (getAllRefs ns xs) t

341 |     getAllRefs ns [] = ns

343 |     updateApp : Name -> RawImp -> RawImp

344 |     updateApp n (IApp fc f a) = IApp fc (updateApp n f) a

345 |     updateApp n (IAutoApp fc f a) = IAutoApp fc (updateApp n f) a

346 |     updateApp n (INamedApp fc f m a) = INamedApp fc (updateApp n f) m a

347 |     updateApp n f = IVar fc n

349 |     unelabDef : (vs ** (Env Term vs, Term vs, Term vs)) ->

350 |                 Core ImpClause

351 |     unelabDef (_ ** (env, lhs, rhs))

352 |         = do lhs' <- unelabNoSugar env lhs

353 |              defs <- get Ctxt

354 |              rhsnf <- normaliseArgHoles defs env rhs

355 |              rhs' <- unelabNoSugar env rhsnf

356 |              pure (PatClause fc (map rawName lhs') (map rawName rhs'))

358 |     showPat : ImpClause -> String

359 |     showPat (PatClause _ lhs rhs) = show lhs ++ " = " ++ show rhs

360 |     showPat _ = "Can't happen"

362 | eraseInferred : {auto c : Ref Ctxt Defs} ->

363 |                 Term vars -> Core (Term vars)

364 | eraseInferred (Bind fc x b tm)

365 |     = do b' <- traverse eraseInferred b

366 |          tm' <- eraseInferred tm

367 |          pure (Bind fc x b' tm')

368 | eraseInferred tm

369 |     = case getFnArgs tm of

370 |            (f, []) => pure f

371 |            (Ref fc Func n, args) =>

372 |                 do defs <- get Ctxt

373 |                    Just gdef <- lookupCtxtExact n (gamma defs)

374 |                         | Nothing => pure tm

375 |                    let argsE = NatSet.overwrite (Erased fc Placeholder) (inferrable gdef) args

376 |                    argsE' <- traverse eraseInferred argsE

377 |                    pure (apply fc (Ref fc Func n) argsE')

378 |            (f, args) =>

379 |                 do args' <- traverse eraseInferred args

380 |                    pure (apply (getLoc f) f args)

385 | specialise : {vars : _} ->

386 |              {auto c : Ref Ctxt Defs} ->

387 |              {auto m : Ref MD Metadata} ->

388 |              {auto u : Ref UST UState} ->

389 |              {auto s : Ref Syn SyntaxInfo} ->

390 |              {auto o : Ref ROpts REPLOpts} ->

391 |              FC -> Env Term vars -> GlobalDef ->

392 |              Name -> List (FC, Term vars) ->

393 |              Core (Maybe (Term vars))

394 | specialise {vars} fc env gdef fn stk

395 |     = let specs = specArgs gdef in

396 |       if NatSet.isEmpty specs then pure Nothing else do

397 |         fnfull <- toFullNames fn

398 |         -- If all the arguments are concrete (meaning, no local variables

399 |         -- or holes in them, so they can be a ClosedTerm) we can specialise

400 |         Just sargs <- getSpecArgs 0 specs stk

401 |             | Nothing => pure Nothing

402 |         defs <- get Ctxt

403 |         sargs <- for sargs $ traversePair $ traverseArgMode $ \ tm =>

404 |                    normalise defs Env.empty tm

405 |         let nhash = hash !(traverse toFullNames $ mapMaybe getStatic $ map snd sargs)

406 |                        `hashWithSalt` fnfull -- add function name to hash to avoid namespace clashes

407 |         let pename = NS partialEvalNS

408 |                      (UN $ Basic ("PE_" ++ nameRoot fnfull ++ "_" ++ asHex (cast nhash)))

409 |         defs <- get Ctxt

410 |         case lookup pename (peFailures defs) of

411 |              Nothing => Just <$> mkSpecDef fc gdef pename sargs fn stk

412 |              Just _ => pure Nothing

413 |   where

414 |     concrete : {vars : _} ->

415 |                Term vars -> Maybe ClosedTerm

416 |     concrete tm = shrink tm none

418 |     getSpecArgs : Nat -> NatSet -> List (FC, Term vars) ->

419 |                   Core (Maybe (List (Nat, ArgMode)))

420 |     getSpecArgs i specs [] = pure (Just [])

421 |     getSpecArgs i specs ((_, x) :: xs)

422 |         = do Just xs' <- getSpecArgs (1 + i) specs xs

423 |                  | Nothing => pure Nothing

424 |              if i `elem` specs

425 |                 then do defs <- get Ctxt

426 |                         x' <- normaliseHoles defs env x

427 |                         x' <- eraseInferred x'

428 |                         let Just xok = concrete x'

429 |                             | Nothing => pure Nothing

430 |                         pure $ Just ((i, Static xok) :: xs')

431 |                 else pure $ Just ((i, Dynamic) :: xs')

433 | findSpecs : {vars : _} ->

434 |             {auto c : Ref Ctxt Defs} ->

435 |             {auto m : Ref MD Metadata} ->

436 |             {auto u : Ref UST UState} ->

437 |             {auto s : Ref Syn SyntaxInfo} ->

438 |             {auto o : Ref ROpts REPLOpts} ->

439 |             Env Term vars -> List (FC, Term vars) -> Term vars ->

440 |             Core (Term vars)

441 | findSpecs env stk (Ref fc Func fn)

442 |     = do defs <- get Ctxt

443 |          Just gdef <- lookupCtxtExact fn (gamma defs)

444 |               | Nothing => pure (applyStackWithFC (Ref fc Func fn) stk)

445 |          Just r <- specialise fc env gdef fn stk

446 |               | Nothing => pure (applyStackWithFC (Ref fc Func fn) stk)

447 |          pure r

448 | findSpecs env stk (Meta fc n i args)

449 |     = do args' <- traverse (findSpecs env []) args

450 |          pure $ applyStackWithFC (Meta fc n i args') stk

451 | findSpecs env stk (Bind fc x b sc)

452 |     = do b' <- traverse (findSpecs env []) b

453 |          sc' <- findSpecs (b' :: env) [] sc

454 |          pure $ applyStackWithFC (Bind fc x b' sc') stk

455 | findSpecs env stk (App fc fn arg)

456 |     = do arg' <- findSpecs env [] arg

457 |          findSpecs env ((fc, arg') :: stk) fn

458 | findSpecs env stk (TDelayed fc r tm)

459 |     = do tm' <- findSpecs env [] tm

460 |          pure $ applyStackWithFC (TDelayed fc r tm') stk

461 | findSpecs env stk (TDelay fc r ty tm)

462 |     = do ty' <- findSpecs env [] ty

463 |          tm' <- findSpecs env [] tm

464 |          pure $ applyStackWithFC (TDelay fc r ty' tm') stk

465 | findSpecs env stk (TForce fc r tm)

466 |     = do tm' <- findSpecs env [] tm

467 |          pure $ applyStackWithFC (TForce fc r tm') stk

468 | findSpecs env stk tm = pure $ applyStackWithFC tm stk

470 | bName : {auto q : Ref QVar Int} -> String -> Core Name

471 | bName n

472 |     = do i <- get QVar

473 |          put QVar (i + 1)

474 |          pure (MN n i)

483 |   quoteArgs : {bound, free : _} ->

484 |               {auto c : Ref Ctxt Defs} ->

485 |               {auto m : Ref MD Metadata} ->

486 |               {auto u : Ref UST UState} ->

487 |               {auto s : Ref Syn SyntaxInfo} ->

488 |               {auto o : Ref ROpts REPLOpts} ->

489 |               Ref QVar Int -> Defs -> Bounds bound ->

490 |               Env Term free -> List (Closure free) ->

491 |               Core (List (Term (bound ++ free)))

492 |   quoteArgs q defs bounds env [] = pure []

493 |   quoteArgs q defs bounds env (a :: args)

494 |       = pure $ (!(quoteGenNF q defs bounds env !(evalClosure defs a)) ::

495 |                 !(quoteArgs q defs bounds env args))

497 |   quoteArgsWithFC : {auto c : Ref Ctxt Defs} ->

498 |                     {auto m : Ref MD Metadata} ->

499 |                     {auto u : Ref UST UState} ->

500 |                     {auto s : Ref Syn SyntaxInfo} ->

501 |                     {auto o : Ref ROpts REPLOpts} ->

502 |                     {bound, free : _} ->

503 |                     Ref QVar Int -> Defs -> Bounds bound ->

504 |                     Env Term free -> List (FC, Closure free) ->

505 |                     Core (List (FC, Term (bound ++ free)))

506 |   quoteArgsWithFC q defs bounds env terms

507 |       = pure $ zip (map fst terms) !(quoteArgs q defs bounds env (map snd terms))

509 |   quoteHead : {bound, free : _} ->

510 |               {auto c : Ref Ctxt Defs} ->

511 |               {auto m : Ref MD Metadata} ->

512 |               {auto u : Ref UST UState} ->

513 |               {auto s : Ref Syn SyntaxInfo} ->

514 |               {auto o : Ref ROpts REPLOpts} ->

515 |               Ref QVar Int -> Defs ->

516 |               FC -> Bounds bound -> Env Term free -> NHead free ->

517 |               Core (Term (bound ++ free))

518 |   quoteHead {bound} q defs fc bounds env (NLocal mrig _ prf)

519 |       = let MkVar prf' = addLater bound prf in

520 |             pure $ Local fc mrig _ prf'

521 |     where

522 |       addLater : {idx : _} -> (ys : List Name) -> (0 p : IsVar n idx xs) ->

523 |                  Var (ys ++ xs)

524 |       addLater [] isv = MkVar isv

525 |       addLater (x :: xs) isv

526 |           = let MkVar isv' = addLater xs isv in

527 |                 MkVar (Later isv')

528 |   quoteHead q defs fc bounds env (NRef Bound (MN n i))

529 |       = case findName bounds of

530 |              Just (MkVar p) => pure $ Local fc Nothing _ (embedIsVar p)

531 |              Nothing => pure $ Ref fc Bound (MN n i)

532 |     where

533 |       findName : Bounds bound' -> Maybe (Var bound')

534 |       findName None = Nothing

535 |       findName (Add x (MN n' i') ns)

536 |           = if i == i' -- this uniquely identifies it, given how we

537 |                        -- generated the names, and is a faster test!

538 |                then Just first

539 |                else do MkVar p <-findName ns

540 |                        Just (MkVar (Later p))

541 |       findName (Add x _ ns)

542 |           = do MkVar p <-findName ns

543 |                Just (MkVar (Later p))

544 |   quoteHead q defs fc bounds env (NRef nt n) = pure $ Ref fc nt n

545 |   quoteHead q defs fc bounds env (NMeta n i args)

546 |       = do args' <- quoteArgs q defs bounds env args

547 |            pure $ Meta fc n i args'

549 |   quotePi : {bound, free : _} ->

550 |             {auto c : Ref Ctxt Defs} ->

551 |             {auto m : Ref MD Metadata} ->

552 |             {auto u : Ref UST UState} ->

553 |             {auto s : Ref Syn SyntaxInfo} ->

554 |             {auto o : Ref ROpts REPLOpts} ->

555 |             Ref QVar Int -> Defs -> Bounds bound ->

556 |             Env Term free -> PiInfo (Closure free) ->

557 |             Core (PiInfo (Term (bound ++ free)))

558 |   quotePi q defs bounds env Explicit = pure Explicit

559 |   quotePi q defs bounds env Implicit = pure Implicit

560 |   quotePi q defs bounds env AutoImplicit = pure AutoImplicit

561 |   quotePi q defs bounds env (DefImplicit t)

562 |       = do t' <- quoteGenNF q defs bounds env !(evalClosure defs t)

563 |            pure (DefImplicit t')

565 |   quoteBinder : {bound, free : _} ->

566 |                 {auto c : Ref Ctxt Defs} ->

567 |                 {auto m : Ref MD Metadata} ->

568 |                 {auto u : Ref UST UState} ->

569 |                 {auto s : Ref Syn SyntaxInfo} ->

570 |                 {auto o : Ref ROpts REPLOpts} ->

571 |                 Ref QVar Int -> Defs -> Bounds bound ->

572 |                 Env Term free -> Binder (Closure free) ->

573 |                 Core (Binder (Term (bound ++ free)))

574 |   quoteBinder q defs bounds env (Lam fc r p ty)

575 |       = do ty' <- quoteGenNF q defs bounds env !(evalClosure defs ty)

576 |            p' <- quotePi q defs bounds env p

577 |            pure (Lam fc r p' ty')

578 |   quoteBinder q defs bounds env (Let fc r val ty)

579 |       = do val' <- quoteGenNF q defs bounds env !(evalClosure defs val)

580 |            ty' <- quoteGenNF q defs bounds env !(evalClosure defs ty)

581 |            pure (Let fc r val' ty')

582 |   quoteBinder q defs bounds env (Pi fc r p ty)

583 |       = do ty' <- quoteGenNF q defs bounds env !(evalClosure defs ty)

584 |            p' <- quotePi q defs bounds env p

585 |            pure (Pi fc r p' ty')

586 |   quoteBinder q defs bounds env (PVar fc r p ty)

587 |       = do ty' <- quoteGenNF q defs bounds env !(evalClosure defs ty)

588 |            p' <- quotePi q defs bounds env p

589 |            pure (PVar fc r p' ty')

590 |   quoteBinder q defs bounds env (PLet fc r val ty)

591 |       = do val' <- quoteGenNF q defs bounds env !(evalClosure defs val)

592 |            ty' <- quoteGenNF q defs bounds env !(evalClosure defs ty)

593 |            pure (PLet fc r val' ty')

594 |   quoteBinder q defs bounds env (PVTy fc r ty)

595 |       = do ty' <- quoteGenNF q defs bounds env !(evalClosure defs ty)

596 |            pure (PVTy fc r ty')

598 |   quoteGenNF : {bound, vars : _} ->

599 |                {auto c : Ref Ctxt Defs} ->

600 |                {auto m : Ref MD Metadata} ->

601 |                {auto u : Ref UST UState} ->

602 |                {auto s : Ref Syn SyntaxInfo} ->

603 |                {auto o : Ref ROpts REPLOpts} ->

604 |                Ref QVar Int ->

605 |                Defs -> Bounds bound ->

606 |                Env Term vars -> NF vars -> Core (Term (bound ++ vars))

607 |   quoteGenNF q defs bound env (NBind fc n b sc)

608 |       = do var <- bName "qv"

609 |            sc' <- quoteGenNF q defs (Add n var bound) env

610 |                        !(sc defs (toClosure defaultOpts env (Ref fc Bound var)))

611 |            b' <- quoteBinder q defs bound env b

612 |            pure (Bind fc n b' sc')

613 |   -- IMPORTANT CASE HERE

614 |   -- If fn is to be specialised, quote the args directly (no further

615 |   -- reduction) then call specialise. Otherwise, quote as normal

616 |   quoteGenNF q defs bound env (NApp fc (NRef Func fn) args)

617 |       = do Just gdef <- lookupCtxtExact fn (gamma defs)

618 |                 | Nothing => do args' <- quoteArgsWithFC q defs bound env args

619 |                                 pure $ applyStackWithFC (Ref fc Func fn) args'

620 |            args' <- quoteArgsWithFC q defs bound env args

621 |            let False = NatSet.isEmpty (specArgs gdef)

622 |                | _ => pure $ applyStackWithFC (Ref fc Func fn) args'

623 |            Just r <- specialise fc (extendEnv bound env) gdef fn args'

624 |                 | Nothing =>

625 |                      -- can't specialise, keep the arguments

626 |                      -- unreduced

627 |                      do empty <- clearDefs defs

628 |                         args' <- quoteArgsWithFC q empty bound env args

629 |                         pure $ applyStackWithFC (Ref fc Func fn) args'

630 |            pure r

631 |      where

632 |        extendEnv : Bounds bs -> Env Term vs -> Env Term (bs ++ vs)

633 |        extendEnv None env = env

634 |        extendEnv (Add x n bs) env

635 |            -- We're just using this to evaluate holes in the right scope, so

636 |            -- a placeholder binder is fine

637 |            = Lam fc top Explicit (Erased fc Placeholder) :: extendEnv bs env

638 |   quoteGenNF q defs bound env (NApp fc f args)

639 |       = do f' <- quoteHead q defs fc bound env f

640 |            args' <- quoteArgsWithFC q defs bound env args

641 |            pure $ applyStackWithFC f' args'

642 |   quoteGenNF q defs bound env (NDCon fc n t ar args)

643 |       = do args' <- quoteArgsWithFC q defs bound env args

644 |            pure $ applyStackWithFC (Ref fc (DataCon t ar) n) args'

645 |   quoteGenNF q defs bound env (NTCon fc n ar args)

646 |       = do args' <- quoteArgsWithFC q defs bound env args

647 |            pure $ applyStackWithFC (Ref fc (TyCon ar) n) args'

648 |   quoteGenNF q defs bound env (NAs fc s n pat)

649 |       = do n' <- quoteGenNF q defs bound env n

650 |            pat' <- quoteGenNF q defs bound env pat

651 |            pure (As fc s n' pat')

652 |   quoteGenNF q defs bound env (NDelayed fc r arg)

653 |       = do argQ <- quoteGenNF q defs bound env arg

654 |            pure (TDelayed fc r argQ)

655 |   quoteGenNF q defs bound env (NDelay fc r ty arg)

656 |       -- unlike main evaluator, we want to look under Delays

657 |       = do argNF <- evalClosure defs arg

658 |            argQ <- quoteGenNF q defs bound env argNF

659 |            tyNF <- evalClosure defs ty

660 |            tyQ <- quoteGenNF q defs bound env tyNF

661 |            pure (TDelay fc r tyQ argQ)

662 |     where

663 |       toHolesOnly : Closure vs -> Closure vs

664 |       toHolesOnly (MkClosure _ locs env tm)

665 |           = MkClosure withHoles locs env tm

666 |       toHolesOnly c = c

667 |   quoteGenNF q defs bound env (NForce fc r arg args)

668 |       = do args' <- quoteArgsWithFC q defs bound env args

669 |            case arg of

670 |                 NDelay fc _ _ arg =>

671 |                    do argNF <- evalClosure defs arg

672 |                       pure $ applyStackWithFC !(quoteGenNF q defs bound env argNF) args'

673 |                 _ => do arg' <- quoteGenNF q defs bound env arg

674 |                         pure $ applyStackWithFC (TForce fc r arg') args'

675 |   quoteGenNF q defs bound env (NPrimVal fc c) = pure $ PrimVal fc c

676 |   quoteGenNF q defs bound env (NErased fc Impossible) = pure $ Erased fc Impossible

677 |   quoteGenNF q defs bound env (NErased fc Placeholder) = pure $ Erased fc Placeholder

678 |   quoteGenNF q defs bound env (NErased fc (Dotted t))

679 |     = pure $ Erased fc $ Dotted !(quoteGenNF q defs bound env t)

680 |   quoteGenNF q defs bound env (NType fc u) = pure $ TType fc u

682 | evalRHS : {vars : _} ->

683 |           {auto c : Ref Ctxt Defs} ->

684 |           {auto m : Ref MD Metadata} ->

685 |           {auto u : Ref UST UState} ->

686 |           {auto s : Ref Syn SyntaxInfo} ->

687 |           {auto o : Ref ROpts REPLOpts} ->

688 |           Env Term vars -> NF vars -> Core (Term vars)

689 | evalRHS env nf

690 |     = do q <- newRef QVar 0

691 |          defs <- get Ctxt

692 |          quoteGenNF q defs None env nf

695 | applySpecialise : {vars : _} ->

696 |                   {auto c : Ref Ctxt Defs} ->

697 |                   {auto m : Ref MD Metadata} ->

698 |                   {auto u : Ref UST UState} ->

699 |                   {auto s : Ref Syn SyntaxInfo} ->

700 |                   {auto o : Ref ROpts REPLOpts} ->

701 |                   Env Term vars ->

702 |                   Maybe (List (Name, Nat)) ->

703 |                         -- ^ If we're specialising, names to reduce in the RHS

704 |                         -- with their reduction limits

705 |                   Term vars -> -- initial RHS

706 |                   Core (Term vars)

707 | applySpecialise env Nothing tm

708 |     = findSpecs env [] tm -- not specialising, just search through RHS

709 | applySpecialise env (Just ls) tmin -- specialising, evaluate RHS while looking

710 |                                  -- for names to specialise

711 |     = do defs <- get Ctxt

712 |          tm <- toResolvedNames tmin

713 |          nf <- nf defs env tm

714 |          tm' <- evalRHS env nf

715 |          tmfull <- toFullNames tm'

716 |          logTermNF "specialise" 5 ("New RHS") env tmfull

717 |          pure tmfull