0 | module TTImp.Elab.Case

  2 | import Core.Env

  3 | import Core.Metadata

  4 | import Core.Unify

  5 | import Core.Value

  7 | import Idris.Syntax

  8 | import Idris.REPL.Opts

 10 | import TTImp.Elab.Check

 11 | import TTImp.Elab.Delayed

 12 | import TTImp.Elab.ImplicitBind

 13 | import TTImp.Elab.Utils

 14 | import TTImp.ProcessFnOpt

 15 | import TTImp.TTImp

 16 | import TTImp.Utils

 18 | import Data.Maybe

 19 | import Data.String

 20 | import Libraries.Data.NameMap

 21 | import Libraries.Data.NatSet

 22 | import Libraries.Data.VarSet

 23 | import Libraries.Data.WithDefault

 25 | %default covering

 28 | changeVar : (old : Var vs) -> (new : Var vs) -> Term vs -> Term vs

 29 | changeVar old (MkVar new) (Local fc r idx p)

 30 |     = if old == MkVar p

 31 |          then Local fc r _ new

 32 |          else Local fc r _ p

 33 | changeVar old new (Meta fc nm i args)

 34 |     = Meta fc nm i (map (changeVar old new) args)

 35 | changeVar (MkVar old) (MkVar new) (Bind fc x b sc)

 36 |     = Bind fc x (assert_total (map (changeVar (MkVar old) (MkVar new)) b))

 37 |            (changeVar (MkVar (Later old)) (MkVar (Later new)) sc)

 38 | changeVar old new (App fc fn arg)

 39 |     = App fc (changeVar old new fn) (changeVar old new arg)

 40 | changeVar old new (As fc s nm p)

 41 |     = As fc s (changeVar old new nm) (changeVar old new p)

 42 | changeVar old new (TDelayed fc r p)

 43 |     = TDelayed fc r (changeVar old new p)

 44 | changeVar old new (TDelay fc r t p)

 45 |     = TDelay fc r (changeVar old new t) (changeVar old new p)

 46 | changeVar old new (TForce fc r p)

 47 |     = TForce fc r (changeVar old new p)

 48 | changeVar old new tm = tm

 50 | toRig1 : {idx : Nat} -> (0 p : IsVar nm idx vs) -> Env Term vs -> Env Term vs

 51 | toRig1 First (b :: bs)

 52 |     = if isErased (multiplicity b)

 53 |          then setMultiplicity b linear :: bs

 54 |          else b :: bs

 55 | toRig1 (Later p) (b :: bs) = b :: toRig1 p bs

 57 | allow : Maybe (Var vs) -> Env Term vs -> Env Term vs

 58 | allow Nothing env = env

 59 | allow (Just (MkVar p)) env = toRig1 p env

 63 | updateMults : VarSet vs -> Env Term vs -> Env Term vs

 64 | updateMults vars env

 65 |   = -- shortcircuiting the call if the set of vars to erase is now empty

 66 |     if VarSet.isEmpty vars then env else go vars env

 67 |   where

 69 |   go : {0 vs : Scope} -> VarSet vs -> Env Term vs -> Env Term vs

 70 |   go vars [] = []

 71 |   go vars (b :: env)

 72 |     = (if first `VarSet.elem` vars

 73 |         then setMultiplicity b erased

 74 |         else b)

 75 |     :: updateMults (VarSet.dropFirst vars) env

 77 | findImpsIn : {vars : _} ->

 78 |              FC -> Env Term vars -> List (Name, Term vars) -> Term vars ->

 79 |              Core ()

 80 | findImpsIn fc env ns (Bind _ n b@(Pi _ _ Implicit ty) sc)

 81 |     = findImpsIn fc (b :: env)

 82 |                  ((n, weaken ty) :: map (\x => (fst x, weaken (snd x))) ns)

 83 |                  sc

 84 | findImpsIn fc env ns (Bind _ n b sc)

 85 |     = findImpsIn fc (b :: env)

 86 |                  (map (\x => (fst x, weaken (snd x))) ns)

 87 |                  sc

 88 | findImpsIn fc env ns ty

 89 |     = when (not (isNil ns)) $

 90 |            throw (TryWithImplicits fc env (reverse ns))

 94 | extendNeeded : {vs : _} ->

 95 |                Binder (Term vs) ->

 96 |                Env Term vs -> VarSet vs -> VarSet vs

 97 | extendNeeded (Let _ _ ty val) env needed

 98 |     = VarSet.union (findUsedLocs env ty) (VarSet.union (findUsedLocs env val) needed)

 99 | extendNeeded (PLet _ _ ty val) env needed

100 |     = VarSet.union (findUsedLocs env ty) (VarSet.union (findUsedLocs env val) needed)

101 | extendNeeded b env needed

102 |     = VarSet.union (findUsedLocs env (binderType b)) needed

104 | findScrutinee : {vs : _} ->

105 |                 Env Term vs -> RawImp -> Maybe (Var vs)

106 | findScrutinee {vs = n' :: _} (b :: bs) (IVar loc' n)

107 |     = if n' == n && not (isLet b)

108 |          then Just first

109 |          else do MkVar p <- findScrutinee bs (IVar loc' n)

110 |                  Just (MkVar (Later p))

111 | findScrutinee _ _ = Nothing

113 | getNestData : (Name, (Maybe Name, List (Var vars), a)) ->

114 |               (Name, Maybe Name, List (Var vars))

115 | getNestData (n, (mn, enames, _)) = (n, mn, enames)

117 | bindCaseLocals : FC -> List (Name, Maybe Name, List (Var vars)) ->

118 |                  List (Name, Name)-> RawImp -> RawImp

119 | bindCaseLocals fc [] args rhs = rhs

120 | bindCaseLocals fc ((n, mn, envns) :: rest) argns rhs

121 |     = -- trace ("Case local " ++ show (n,mn,envns) ++ " from " ++ show argns) $

122 |         ICaseLocal fc n (fromMaybe n mn)

123 |                  (map getNameFrom envns)

124 |                  (bindCaseLocals fc rest argns rhs)

125 |   where

126 |     getArg : List (Name, Name) -> Nat -> Maybe Name

127 |     getArg [] _ = Nothing

128 |     getArg ((_, x) :: xs) Z = Just x

129 |     getArg (x :: xs) (S k) = getArg xs k

131 |     getNameFrom : Var vars -> Name

132 |     getNameFrom (MkVar {varIdx = i} _)

133 |         = case getArg argns i of

134 |                Nothing => n

135 |                Just n' => n'

138 | caseBlock : {vars : _} ->

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

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

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

142 |             {auto e : Ref EST (EState vars)} ->

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

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

145 |             RigCount ->

146 |             ElabInfo -> FC ->

147 |             NestedNames vars ->

148 |             Env Term vars ->

149 |             List FnOpt ->

150 |             RawImp -> -- original scrutinee

151 |             Term vars -> -- checked scrutinee

152 |             Term vars -> -- its type

153 |             RigCount -> -- its multiplicity

154 |             List ImpClause -> Maybe (Glued vars) ->

155 |             Core (Term vars, Glued vars)

156 | caseBlock {vars} rigc elabinfo fc nest env opts scr scrtm scrty caseRig alts expected

157 |     = do -- TODO (or to decide): Blodwen allowed ambiguities in the scrutinee

158 |          -- to be delayed, but now I think it's better to have simpler

159 |          -- resolution rules, and not delay

161 |          est <- get EST

162 |          fullImps <- getToBind fc (elabMode elabinfo)

163 |                                (implicitMode elabinfo) env []

164 |          log "elab.case" 5 $ "Doing a case under unbound implicits " ++ show fullImps

166 |          scrn <- genVarName "scr"

167 |          casen <- genCaseName !(prettyName !(toFullNames (Resolved (defining est))))

169 |          -- Update environment so that linear bindings which were used

170 |          -- (esp. in the scrutinee!) are set to 0 in the case type

171 |          let env = updateMults (linearUsed est) env

172 |          defs <- get Ctxt

173 |          parentDef <- lookupCtxtExact (Resolved (defining est)) (gamma defs)

174 |          let vis = specified $ case parentDef of

175 |                         Just gdef =>

176 |                              if collapseDefault (visibility gdef) == Public

177 |                                 then Public

178 |                                 else Private

179 |                         Nothing => Public

181 |          -- if the scrutinee is ones of the arguments in 'env' we should

182 |          -- split on that, rather than adding it as a new argument

183 |          let splitOn = findScrutinee env scr

185 |          caseretty_in <- case expected of

186 |                            Just ty => getTerm ty

187 |                            _ =>

188 |                               do nmty <- genName "caseTy"

189 |                                  u <- uniVar fc

190 |                                  metaVar fc erased env nmty (TType fc u)

192 |          u <- uniVar fc

193 |          (caseretty, _) <- bindImplicits fc (implicitMode elabinfo) defs env

194 |                                          fullImps caseretty_in (TType fc u)

195 |          let casefnty

196 |                = abstractFullEnvType fc (allow splitOn (mkExplicit env))

197 |                             (maybe (Bind fc scrn (Pi fc caseRig Explicit scrty)

198 |                                        (weaken caseretty))

199 |                                    (const caseretty) splitOn)

200 |          -- If we can normalise the type without the result being excessively

201 |          -- big do it. It's the depth of stuck applications - 10 is already

202 |          -- pretty much unreadable!

203 |          casefnty <- normaliseSizeLimit defs 10 Env.empty casefnty

204 |          (erasedargs, _) <- findErased casefnty

206 |          logEnv "elab.case" 10 "Case env" env

207 |          logTermNF "elab.case" 2 ("Case function type: " ++ show casen) Env.empty casefnty

208 |          traverse_ addToSave (keys (getMetas casefnty))

210 |          -- If we've had to add implicits to the case type (because there

211 |          -- were unbound implicits) then we're in a bit of a mess. Easiest

212 |          -- way out is to throw an error and try again with the implicits

213 |          -- actually bound! This is rather hacky, but a lot less fiddly than

214 |          -- the alternative of fixing up the environment

215 |          when (not (isNil fullImps)) $ findImpsIn fc Env.empty [] casefnty

216 |          cidx <- addDef casen ({ eraseArgs := erasedargs }

217 |                                 (newDef fc casen (if isErased rigc then erased else top)

218 |                                       Scope.empty casefnty vis None))

220 |          traverse_ (processFnOpt fc False casen) opts

222 |          -- set the totality of the case block to be the same as that

223 |          -- of the parent function

224 |          let tot = fromMaybe PartialOK $ do findSetTotal (flags !parentDef)

225 |          log "elab.case" 5 $

226 |            unwords [ "Setting totality requirement for", show casen

227 |                    , "to", show tot]

228 |          setFlag fc (Resolved cidx) (SetTotal tot)

229 |          let caseRef : Term vars = Ref fc Func (Resolved cidx)

231 |          let applyEnv = applyToFull fc caseRef env

232 |          let appTm : Term vars

233 |                    = maybe (App fc applyEnv scrtm)

234 |                            (const applyEnv)

235 |                            splitOn

237 |          let alts' = map (updateClause casen splitOn nest env) alts

238 |          log "elab.case" 2 $ "Nested: " ++ show (map getNestData (names nest))

239 |          log "elab.case" 2 $ "Generated alts: " ++ show alts'

240 |          logTermNF "elab.case" 2 "Case application" env appTm

242 |          -- Start with empty nested names, since we've extended the rhs with

243 |          -- ICaseLocal so they'll get rebuilt with the right environment

244 |          let nest' = MkNested []

245 |          ust <- get UST

246 |          -- We don't want to keep rechecking delayed elaborators in the

247 |          -- case block, because they're not going to make progress until

248 |          -- we come out again, so save them

249 |          let olddelayed = delayedElab ust

250 |          put UST ({ delayedElab := [] } ust)

251 |          processDecl [InCase] nest' Env.empty (IDef fc casen alts')

253 |          -- If there's no duplication of the scrutinee in the block,

254 |          -- flag it as inlinable.

255 |          -- This will be the case either if the scrutinee is a variable, in

256 |          -- which case the duplication won't hurt, or if there's no variable

257 |          -- duplicated in the body (what ghc calls W-safe)

258 |          -- We'll check that second condition later, after generating the

259 |          -- runtime (erased) case trees

260 |          let inlineOK = maybe False (const True) splitOn

261 |          when inlineOK $ setFlag fc casen Inline

263 |          ust <- get UST

264 |          put UST ({ delayedElab := olddelayed } ust)

266 |          pure (appTm, gnf env caseretty)

267 |   where

268 |     mkLocalEnv : Env Term vs -> Env Term vs

269 |     mkLocalEnv [] = Env.empty

270 |     mkLocalEnv (b :: bs)

271 |         = let b' = if isLinear (multiplicity b)

272 |                       then setMultiplicity b erased

273 |                       else b in

274 |               b' :: mkLocalEnv bs

276 |     -- Return the original name in the environment, and what it needs to be

277 |     -- called in the case block. We need to mapping to build the ICaseLocal

278 |     -- so that it applies to the right original variable

279 |     getBindName : Int -> Name -> List Name -> (Name, Name)

280 |     getBindName idx n@(UN un) vs

281 |        = if n `elem` vs then (n, MN (displayUserName un) idx) else (n, n)

282 |     getBindName idx n vs

283 |        = if n `elem` vs then (n, MN "_cn" idx) else (n, n)

285 |     -- Returns a list of names that nestednames should be applied to, mapped

286 |     -- to what the name has become in the case block, and a list of terms for

287 |     -- the LHS of the case to be applied to.

288 |     addEnv : {vs : _} ->

289 |              Int -> Env Term vs -> List Name -> (List (Name, Name), List RawImp)

290 |     addEnv idx [] used = ([], [])

291 |     addEnv idx {vs = v :: vs} (b :: bs) used

292 |         = let n = getBindName idx v used

293 |               (ns, rest) = addEnv (idx + 1) bs (snd n :: used)

294 |               ns' = n :: ns in

295 |               (ns', IAs fc EmptyFC UseLeft (snd n) (Implicit fc True) :: rest)

297 |     -- Replace a variable in the argument list; if the reference is to

298 |     -- a variable kept in the outer environment (therefore not an argument

299 |     -- in the list) don't consume it

300 |     replace : (idx : Nat) -> RawImp -> List RawImp -> List RawImp

301 |     replace Z lhs (old :: xs)

302 |        = let lhs' = case old of

303 |                          IAs loc' nameLoc' side n _ => IAs loc' nameLoc' side n lhs

304 |                          _ => lhs in

305 |              lhs' :: xs

306 |     replace (S k) lhs (x :: xs)

307 |         = x :: replace k lhs xs

308 |     replace _ _ xs = xs

310 |     mkSplit : Maybe (Var vs) ->

311 |               RawImp -> List RawImp ->

312 |               List RawImp

313 |     mkSplit Nothing lhs args = reverse (lhs :: args)

314 |     mkSplit (Just (MkVar {varIdx = i} prf)) lhs args

315 |         = reverse (replace i lhs args)

317 |     -- Names used in the pattern we're matching on, so don't bind them

318 |     -- in the generated case block

319 |     usedIn : RawImp -> List Name

320 |     usedIn (IBindVar _ n) = [n]

321 |     usedIn (IApp _ f a) = usedIn f ++ usedIn a

322 |     usedIn (IAs _ _ _ n a) = n :: usedIn a

323 |     usedIn (IAlternative _ _ alts) = concatMap usedIn alts

324 |     usedIn _ = []

326 |     -- Get a name update for the LHS (so that if there's a nested data declaration

327 |     -- the constructors are applied to the environment in the case block)

328 |     nestLHS : FC -> (Name, (Maybe Name, List (Var vars), a)) -> (Name, RawImp)

329 |     nestLHS fc (n, (mn, ns, t))

330 |         = (n, apply (IVar fc (fromMaybe n mn))

331 |                     (map (const (Implicit fc False)) ns))

333 |     applyNested : NestedNames vars -> RawImp -> RawImp

334 |     applyNested nest lhs

335 |         = let fc = getFC lhs in

336 |               substNames [] (map (nestLHS fc) (names nest)) lhs

338 |     updateClause : Name -> Maybe (Var vars) ->

339 |                    NestedNames vars ->

340 |                    Env Term vars -> ImpClause -> ImpClause

341 |     updateClause casen splitOn nest env (PatClause loc' lhs rhs)

342 |         = let (ns, args) = addEnv 0 env (usedIn lhs)

343 |               args' = mkSplit splitOn lhs args

344 |               lhs' = apply (IVar loc' casen) args' in

345 |               PatClause loc' (applyNested nest lhs')

346 |                         (bindCaseLocals loc' (map getNestData (names nest))

347 |                                         ns rhs)

348 |     -- With isn't allowed in a case block but include for completeness

349 |     updateClause casen splitOn nest env (WithClause loc' lhs rig wval prf flags cs)

350 |         = let (_, args) = addEnv 0 env (usedIn lhs)

351 |               args' = mkSplit splitOn lhs args

352 |               lhs' = apply (IVar loc' casen) args' in

353 |               WithClause loc' (applyNested nest lhs') rig wval prf flags cs

354 |     updateClause casen splitOn nest env (ImpossibleClause loc' lhs)

355 |         = let (_, args) = addEnv 0 env (usedIn lhs)

356 |               args' = mkSplit splitOn lhs args

357 |               lhs' = apply (IVar loc' casen) args' in

358 |               ImpossibleClause loc' (applyNested nest lhs')

362 | checkCase : {vars : _} ->

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

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

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

366 |             {auto e : Ref EST (EState vars)} ->

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

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

369 |             RigCount -> ElabInfo ->

370 |             NestedNames vars -> Env Term vars ->

371 |             FC -> List FnOpt -> (scr : RawImp) -> (ty : RawImp) -> List ImpClause ->

372 |             Maybe (Glued vars) ->

373 |             Core (Term vars, Glued vars)

374 | checkCase rig elabinfo nest env fc opts scr scrty_in alts exp

375 |     = delayElab fc rig env exp CaseBlock $

376 |         do scrty_exp <- case scrty_in of

377 |                              Implicit {} => guessScrType alts

378 |                              _ => pure scrty_in

379 |            u <- uniVar fc

380 |            (scrtyv, scrtyt) <- check erased elabinfo nest env scrty_exp

381 |                                      (Just (gType fc u))

382 |            logTerm "elab.case" 10 "Expected scrutinee type" scrtyv

383 |            -- Try checking at the given multiplicity; if that doesn't work,

384 |            -- try checking at Rig1 (meaning that we're using a linear variable

385 |            -- so the scrutinee should be linear)

386 |            let chrig = if isErased rig then erased else top

387 |            log "elab.case" 5 $ "Checking " ++ show scr ++ " at " ++ show chrig

389 |            (scrtm_in, gscrty, caseRig) <- handle

390 |               (do c <- runDelays (const True) $ check chrig elabinfo nest env scr (Just (gnf env scrtyv))

391 |                   pure (fst c, snd c, chrig))

392 |             $ \case

393 |                 e@(LinearMisuse _ _ r _)

394 |                   => branchOne

395 |                      (do c <- runDelays (const True) $ check linear elabinfo nest env scr

396 |                               (Just (gnf env scrtyv))

397 |                          pure (fst c, snd c, linear))

398 |                      (throw e)

399 |                      r

400 |                 e => throw e

402 |            scrty <- getTerm gscrty

403 |            logTermNF "elab.case" 5 "Scrutinee type" env scrty

404 |            defs <- get Ctxt

405 |            checkConcrete !(nf defs env scrty)

406 |            caseBlock rig elabinfo fc nest env opts scr scrtm_in scrty caseRig alts exp

407 |   where

408 |     -- For the moment, throw an error if we haven't been able to work out

409 |     -- the type of the case scrutinee, because we'll need it to build the

410 |     -- type of the case block. But (TODO) consider delaying on failure?

411 |     checkConcrete : NF vs -> Core ()

412 |     checkConcrete (NApp _ (NMeta n i _) _)

413 |         = throw (GenericMsg fc "Can't infer type for case scrutinee")

414 |     checkConcrete _ = pure ()

416 |     applyTo : Defs -> RawImp -> ClosedNF -> Core RawImp

417 |     applyTo defs ty (NBind fc _ (Pi _ _ Explicit _) sc)

418 |         = applyTo defs (IApp fc ty (Implicit fc False))

419 |                !(sc defs (toClosure defaultOpts Env.empty (Erased fc Placeholder)))

420 |     applyTo defs ty (NBind _ x (Pi {}) sc)

421 |         = applyTo defs (INamedApp fc ty x (Implicit fc False))

422 |                !(sc defs (toClosure defaultOpts Env.empty (Erased fc Placeholder)))

423 |     applyTo defs ty _ = pure ty

425 |     -- Get the name and type of the family the scrutinee is in

426 |     getRetTy : Defs -> ClosedNF -> Core (Maybe (Name, ClosedNF))

427 |     getRetTy defs (NBind fc _ (Pi {}) sc)

428 |         = getRetTy defs !(sc defs (toClosure defaultOpts Env.empty (Erased fc Placeholder)))

429 |     getRetTy defs (NTCon _ n arity _)

430 |         = do Just ty <- lookupTyExact n (gamma defs)

431 |                   | Nothing => pure Nothing

432 |              pure (Just (n, !(nf defs Env.empty ty)))

433 |     getRetTy _ _ = pure Nothing

435 |     -- Guess a scrutinee type by looking at the alternatives, so that we

436 |     -- have a hint for building the case type

437 |     guessScrType : List ImpClause -> Core RawImp

438 |     guessScrType [] = pure $ Implicit fc False

439 |     guessScrType (PatClause _ x _ :: xs)

440 |         = case getFn x of

441 |                IVar _ n =>

442 |                   do defs <- get Ctxt

443 |                      [(_, (_, ty))] <- lookupTyName (mapNestedName nest n) (gamma defs)

444 |                          | _ => guessScrType xs

445 |                      Just (tyn, tyty) <- getRetTy defs !(nf defs Env.empty ty)

446 |                          | _ => guessScrType xs

447 |                      applyTo defs (IVar fc tyn) tyty

448 |                _ => guessScrType xs

449 |     guessScrType (_ :: xs) = guessScrType xs