Skip to content

Commit

Permalink
chore: simplify signature of Array.mapIdx
Browse files Browse the repository at this point in the history
  • Loading branch information
kim-em committed Oct 17, 2024
1 parent 3a34a8e commit c71b15a
Show file tree
Hide file tree
Showing 10 changed files with 27 additions and 29 deletions.
9 changes: 4 additions & 5 deletions src/Init/Data/Array/Basic.lean
Original file line number Diff line number Diff line change
Expand Up @@ -397,17 +397,16 @@ def mapM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (f : α
map 0 (mkEmpty as.size)

@[inline]
def mapIdxM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : Array α) (f : Fin as.size → α → m β) : m (Array β) :=
def mapIdxM {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : Array α) (f : Nat → α → m β) : m (Array β) :=
let rec @[specialize] map (i : Nat) (j : Nat) (inv : i + j = as.size) (bs : Array β) : m (Array β) := do
match i, inv with
| 0, _ => pure bs
| i+1, inv =>
have : j < as.size := by
have j_lt : j < as.size := by
rw [← inv, Nat.add_assoc, Nat.add_comm 1 j, Nat.add_comm]
apply Nat.le_add_right
let idx : Fin as.size := ⟨j, this⟩
have : i + (j + 1) = as.size := by rw [← inv, Nat.add_comm j 1, Nat.add_assoc]
map i (j+1) this (bs.push (← f idx (as.get idx)))
map i (j+1) this (bs.push (← f j (as.get ⟨j, j_lt⟩)))
map as.size 0 rfl (mkEmpty as.size)

@[inline]
Expand Down Expand Up @@ -516,7 +515,7 @@ def map {α : Type u} {β : Type v} (f : α → β) (as : Array α) : Array β :
Id.run <| as.mapM f

@[inline]
def mapIdx {α : Type u} {β : Type v} (as : Array α) (f : Fin as.size → α → β) : Array β :=
def mapIdx {α : Type u} {β : Type v} (as : Array α) (f : Nat → α → β) : Array β :=
Id.run <| as.mapIdxM f

/-- Turns `#[a, b]` into `#[(a, 0), (b, 1)]`. -/
Expand Down
19 changes: 9 additions & 10 deletions src/Init/Data/Array/MapIdx.lean
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ namespace Array
/-! ### mapIdx -/

-- This could also be proved from `SatisfiesM_mapIdxM` in Batteries.
theorem mapIdx_induction (as : Array α) (f : Fin as.size → α → β)
theorem mapIdx_induction (as : Array α) (f : Nat → α → β)
(motive : Nat → Prop) (h0 : motive 0)
(p : Fin as.size → β → Prop)
(hs : ∀ i, motive i.1 → p i (f i as[i]) ∧ motive (i + 1)) :
Expand All @@ -27,7 +27,7 @@ theorem mapIdx_induction (as : Array α) (f : Fin as.size → α → β)
have := (Nat.zero_add _).symm.trans h
exact ⟨this ▸ hm, h₁ ▸ this, fun _ _ => h₂ ..⟩
| succ i ih =>
apply @ih (bs.push (f ⟨j, by omega⟩ as[j])) (j + 1) (by omega) (by simp; omega)
apply @ih (bs.push (f j as[j])) (j + 1) (by omega) (by simp; omega)
· intro i i_lt h'
rw [get_push]
split
Expand All @@ -38,26 +38,25 @@ theorem mapIdx_induction (as : Array α) (f : Fin as.size → α → β)
· exact (hs ⟨j, by omega⟩ hm).2
simp [mapIdx, mapIdxM]; exact go rfl nofun h0

theorem mapIdx_spec (as : Array α) (f : Fin as.size → α → β)
theorem mapIdx_spec (as : Array α) (f : Nat → α → β)
(p : Fin as.size → β → Prop) (hs : ∀ i, p i (f i as[i])) :
∃ eq : (Array.mapIdx as f).size = as.size,
∀ i h, p ⟨i, h⟩ ((Array.mapIdx as f)[i]) :=
(mapIdx_induction _ _ (fun _ => True) trivial p fun _ _ => ⟨hs .., trivial⟩).2

@[simp] theorem size_mapIdx (a : Array α) (f : Fin a.size → α → β) : (a.mapIdx f).size = a.size :=
@[simp] theorem size_mapIdx (a : Array α) (f : Nat → α → β) : (a.mapIdx f).size = a.size :=
(mapIdx_spec (p := fun _ _ => True) (hs := fun _ => trivial)).1

@[simp] theorem size_zipWithIndex (as : Array α) : as.zipWithIndex.size = as.size :=
Array.size_mapIdx _ _

@[simp] theorem getElem_mapIdx (a : Array α) (f : Fin a.size → α → β) (i : Nat)
@[simp] theorem getElem_mapIdx (a : Array α) (f : Nat → α → β) (i : Nat)
(h : i < (mapIdx a f).size) :
(a.mapIdx f)[i] = f ⟨i, by simp_all⟩ (a[i]'(by simp_all)) :=
(mapIdx_spec _ _ (fun i b => b = f i a[i]) fun _ => rfl).2 i _
(a.mapIdx f)[i] = f i (a[i]'(by simp_all)) :=
(mapIdx_spec _ _ (fun i b => b = f i a[i]) fun _ => rfl).2 i (by simpa using h)

@[simp] theorem getElem?_mapIdx (a : Array α) (f : Fin a.size → α → β) (i : Nat) :
(a.mapIdx f)[i]? =
a[i]?.pbind fun b h => f ⟨i, (getElem?_eq_some_iff.1 h).1⟩ b := by
@[simp] theorem getElem?_mapIdx (a : Array α) (f : Nat → α → β) (i : Nat) :
(a.mapIdx f)[i]? = a[i]?.map (f i) := by
simp only [getElem?_def, size_mapIdx, getElem_mapIdx]
split <;> simp_all

Expand Down
2 changes: 1 addition & 1 deletion src/Lean/Elab/Syntax.lean
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ where
let args ← args.mapM fun arg => withNestedParser do process arg
mkParserSeq args
else
let args ← args.mapIdxM fun i arg => withReader (fun ctx => { ctx with first := ctx.first && i.val == 0 }) do process arg
let args ← args.mapIdxM fun i arg => withReader (fun ctx => { ctx with first := ctx.first && i == 0 }) do process arg
mkParserSeq args

ensureNoPrec (stx : Syntax) :=
Expand Down
2 changes: 1 addition & 1 deletion src/Lean/Elab/Tactic/Conv/Pattern.lean
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ private def pre (pattern : AbstractMVarsResult) (state : IO.Ref PatternMatchStat
let ids ← ids.mapIdxM fun i id =>
match id.getNat with
| 0 => throwErrorAt id "positive integer expected"
| n+1 => pure (n, i.1)
| n+1 => pure (n, i)
let ids := ids.qsort (·.1 < ·.1)
unless @Array.allDiff _ ⟨(·.1 == ·.1)⟩ ids do
throwError "occurrence list is not distinct"
Expand Down
10 changes: 5 additions & 5 deletions src/Lean/Meta/IndPredBelow.lean
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ def mkContext (declName : Name) : MetaM Context := do
let typeInfos ← indVal.all.toArray.mapM getConstInfoInduct
let motiveTypes ← typeInfos.mapM motiveType
let motives ← motiveTypes.mapIdxM fun j motive =>
return (← motiveName motiveTypes j.val, motive)
return (← motiveName motiveTypes j, motive)
let headers ← typeInfos.mapM $ mkHeader motives indVal.numParams
return {
motives := motives
Expand Down Expand Up @@ -214,7 +214,7 @@ def mkConstructor (ctx : Context) (i : Nat) (ctor : Name) : MetaM Constructor :=

def mkInductiveType
(ctx : Context)
(i : Fin ctx.typeInfos.size)
(i : Nat)
(indVal : InductiveVal) : MetaM InductiveType := do
return {
name := ctx.belowNames[i]!
Expand Down Expand Up @@ -340,11 +340,11 @@ where
mkIH
(params : Array Expr)
(motives : Array Expr)
(idx : Fin ctx.motives.size)
(idx : Nat)
(motive : Name × Expr) : MetaM $ Name × (Array Expr → MetaM Expr) := do
let name :=
if ctx.motives.size > 1
then mkFreshUserName <| .mkSimple s!"ih_{idx.val.succ}"
then mkFreshUserName <| .mkSimple s!"ih_{idx + 1}"
else mkFreshUserName <| .mkSimple "ih"
let ih ← instantiateForall motive.2 params
let mkDomain (_ : Array Expr) : MetaM Expr :=
Expand All @@ -353,7 +353,7 @@ where
let args := params ++ motives ++ ys
let premise :=
mkAppN
(mkConst ctx.belowNames[idx.val]! levels) args
(mkConst ctx.belowNames[idx]! levels) args
let conclusion :=
mkAppN motives[idx]! ys
mkForallFVars ys (←mkArrow premise conclusion)
Expand Down
2 changes: 1 addition & 1 deletion src/Lean/Meta/Match/CaseArraySizes.lean
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ def caseArraySizes (mvarId : MVarId) (fvarId : FVarId) (sizes : Array Nat) (xNam
let subst := subgoal.subst
let mvarId := subgoal.mvarId
let hEqSz := (subst.get hEq).fvarId!
if h : i.val < sizes.size then
if h : i < sizes.size then
let n := sizes.get ⟨i, h⟩
let mvarId ← mvarId.clear subgoal.newHs[0]!
let mvarId ← mvarId.clear (subst.get aSizeFVarId).fvarId!
Expand Down
4 changes: 2 additions & 2 deletions src/Lean/Meta/Match/Match.lean
Original file line number Diff line number Diff line change
Expand Up @@ -545,7 +545,7 @@ private def processValue (p : Problem) : MetaM (Array Problem) := do
let subgoals ← caseValues p.mvarId x.fvarId! values (substNewEqs := true)
subgoals.mapIdxM fun i subgoal => do
trace[Meta.Match.match] "processValue subgoal\n{MessageData.ofGoal subgoal.mvarId}"
if h : i.val < values.size then
if h : i < values.size then
let value := values.get ⟨i, h⟩
-- (x = value) branch
let subst := subgoal.subst
Expand Down Expand Up @@ -599,7 +599,7 @@ private def processArrayLit (p : Problem) : MetaM (Array Problem) := do
let sizes := collectArraySizes p
let subgoals ← caseArraySizes p.mvarId x.fvarId! sizes
subgoals.mapIdxM fun i subgoal => do
if i.val < sizes.size then
if i < sizes.size then
let size := sizes.get! i
let subst := subgoal.subst
let elems := subgoal.elems.toList
Expand Down
4 changes: 2 additions & 2 deletions src/Lean/Meta/Tactic/FunInd.lean
Original file line number Diff line number Diff line change
Expand Up @@ -643,7 +643,7 @@ def abstractIndependentMVars (mvars : Array MVarId) (index : Nat) (e : Expr) : M
pure mvar
trace[Meta.FunInd] "abstractIndependentMVars, reverted mvars: {mvars}"
let decls := mvars.mapIdx fun i mvar =>
(.mkSimple s!"case{i.val+1}", (fun _ => mvar.getType))
(.mkSimple s!"case{i+1}", (fun _ => mvar.getType))
Meta.withLocalDeclsD decls fun xs => do
for mvar in mvars, x in xs do
mvar.assign x
Expand Down Expand Up @@ -971,7 +971,7 @@ def deriveInductionStructural (names : Array Name) (numFixed : Nat) : MetaM Unit
mkForallFVars ys (.sort levelZero)
let motiveArities ← infos.mapM fun info => do
lambdaTelescope (← instantiateLambda info.value xs) fun ys _ => pure ys.size
let motiveDecls ← motiveTypes.mapIdxM fun ⟨i,_⟩ motiveType => do
let motiveDecls ← motiveTypes.mapIdxM fun i motiveType => do
let n := if infos.size = 1 then .mkSimple "motive"
else .mkSimple s!"motive_{i+1}"
pure (n, fun _ => pure motiveType)
Expand Down
2 changes: 1 addition & 1 deletion src/Lean/Parser/Extra.lean
Original file line number Diff line number Diff line change
Expand Up @@ -189,7 +189,7 @@ open PrettyPrinter Syntax.MonadTraverser Formatter in
@[combinator_formatter sepByIndent]
def sepByIndent.formatter (p : Formatter) (_sep : String) (pSep : Formatter) : Formatter := do
let stx ← getCur
let hasNewlineSep := stx.getArgs.mapIdx (fun ⟨i, _⟩ n =>
let hasNewlineSep := stx.getArgs.mapIdx (fun i n =>
i % 2 == 1 && n.matchesNull 0 && i != stx.getArgs.size - 1) |>.any id
visitArgs do
for i in (List.range stx.getArgs.size).reverse do
Expand Down
2 changes: 1 addition & 1 deletion src/Lean/Server/Completion.lean
Original file line number Diff line number Diff line change
Expand Up @@ -1004,7 +1004,7 @@ private def assignSortTexts (completions : CompletionList) : CompletionList := I
if completions.items.isEmpty then
return completions
let items := completions.items.mapIdx fun i item =>
{ item with sortText? := toString i.val }
{ item with sortText? := toString i }
let maxDigits := items[items.size - 1]!.sortText?.get!.length
let items := items.map fun item =>
let sortText := item.sortText?.get!
Expand Down

0 comments on commit c71b15a

Please sign in to comment.