Skip to content

Commit

Permalink
feat: simprocs for #[1,2,3,4,5][2] (#4765)
Browse files Browse the repository at this point in the history
None of these were working previously:

```
#check_simp #[1,2,3,4,5][2]  ~> 3
#check_simp #[1,2,3,4,5][2]? ~> some 3
#check_simp #[1,2,3,4,5][7]? ~> none
#check_simp #[][0]? ~> none
#check_simp #[1,2,3,4,5][2]! ~> 3
#check_simp #[1,2,3,4,5][7]! ~> (default : Nat)
#check_simp (#[] : Array Nat)[0]! ~> (default : Nat)
```
  • Loading branch information
kim-em authored Jul 17, 2024
1 parent f6666fe commit af03af5
Show file tree
Hide file tree
Showing 5 changed files with 96 additions and 0 deletions.
8 changes: 8 additions & 0 deletions src/Lean/Meta/AppBuilder.lean
Original file line number Diff line number Diff line change
Expand Up @@ -504,6 +504,14 @@ def mkArrayLit (type : Expr) (xs : List Expr) : MetaM Expr := do
let listLit ← mkListLit type xs
return mkApp (mkApp (mkConst ``List.toArray [u]) type) listLit

def mkNone (type : Expr) : MetaM Expr := do
let u ← getDecLevel type
return mkApp (mkConst ``Option.none [u]) type

def mkSome (type value : Expr) : MetaM Expr := do
let u ← getDecLevel type
return mkApp2 (mkConst ``Option.some [u]) type value

def mkSorry (type : Expr) (synthetic : Bool) : MetaM Expr := do
let u ← getLevel type
return mkApp2 (mkConst ``sorryAx [u]) type (toExpr synthetic)
Expand Down
44 changes: 44 additions & 0 deletions src/Lean/Meta/LitValues.lean
Original file line number Diff line number Diff line change
Expand Up @@ -176,4 +176,48 @@ def litToCtor (e : Expr) : MetaM Expr := do
return mkApp3 (mkConst ``Fin.mk) n i h
return e

/--
Check if an expression is a list literal (i.e. a nested chain of `List.cons`, ending at a `List.nil`),
where each element is "recognised" by a given function `f : Expr → MetaM (Option α)`,
and return the array of recognised values.
-/
partial def getListLitOf? (e : Expr) (f : Expr → MetaM (Option α)) : MetaM (Option (Array α)) := do
let mut e ← instantiateMVars e.consumeMData
let mut r := #[]
while true do
match_expr e with
| List.nil _ => break
| List.cons _ a as => do
let some a ← f a | return none
r := r.push a
e := as
| _ => return none
return some r

/--
Check if an expression is a list literal (i.e. a nested chain of `List.cons`, ending at a `List.nil`),
returning the array of `Expr` values.
-/
def getListLit? (e : Expr) : MetaM (Option (Array Expr)) := getListLitOf? e fun s => return some s

/--
Check if an expression is an array literal
(i.e. `List.toArray` applied to a nested chain of `List.cons`, ending at a `List.nil`),
where each element is "recognised" by a given function `f : Expr → MetaM (Option α)`,
and return the array of recognised values.
-/
def getArrayLitOf? (e : Expr) (f : Expr → MetaM (Option α)) : MetaM (Option (Array α)) := do
let e ← instantiateMVars e.consumeMData
match_expr e with
| List.toArray _ as => getListLitOf? as f
| _ => return none

/--
Check if an expression is an array literal
(i.e. `List.toArray` applied to a nested chain of `List.cons`, ending at a `List.nil`),
returning the array of `Expr` values.
-/
def getArrayLit? (e : Expr) : MetaM (Option (Array Expr)) := getArrayLitOf? e fun s => return some s


end Lean.Meta
1 change: 1 addition & 0 deletions src/Lean/Meta/Tactic/Simp/BuiltinSimprocs.lean
Original file line number Diff line number Diff line change
Expand Up @@ -13,3 +13,4 @@ import Lean.Meta.Tactic.Simp.BuiltinSimprocs.Char
import Lean.Meta.Tactic.Simp.BuiltinSimprocs.String
import Lean.Meta.Tactic.Simp.BuiltinSimprocs.BitVec
import Lean.Meta.Tactic.Simp.BuiltinSimprocs.List
import Lean.Meta.Tactic.Simp.BuiltinSimprocs.Array
36 changes: 36 additions & 0 deletions src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/Array.lean
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
/-
Copyright (c) 2024 Lean FRO. All Rights Reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Kim Morrison
-/
prelude
import Lean.Meta.LitValues
import Lean.Meta.Tactic.Simp.BuiltinSimprocs.Nat

namespace Array
open Lean Meta Simp

/-- Simplification procedure for `#[...][n]` for `n` a `Nat` literal. -/
builtin_dsimproc [simp, seval] reduceGetElem (@GetElem.getElem (Array _) Nat _ _ _ _ _ _) := fun e => do
let_expr GetElem.getElem _ _ _ _ _ xs n _ ← e | return .continue
let some n ← Nat.fromExpr? n | return .continue
let some xs ← getArrayLit? xs | return .continue
return .done <| xs[n]!

/-- Simplification procedure for `#[...][n]?` for `n` a `Nat` literal. -/
builtin_dsimproc [simp, seval] reduceGetElem? (@GetElem?.getElem? (Array _) Nat _ _ _ _ _) := fun e => do
let_expr GetElem?.getElem? _ _ α _ _ xs n ← e | return .continue
let some n ← Nat.fromExpr? n | return .continue
let some xs ← getArrayLit? xs | return .continue
let r ← if h : n < xs.size then mkSome α xs[n] else mkNone α
return .done r

/-- Simplification procedure for `#[...][n]!` for `n` a `Nat` literal. -/
builtin_dsimproc [simp, seval] reduceGetElem! (@GetElem?.getElem! (Array _) Nat _ _ _ _ _ _) := fun e => do
let_expr GetElem?.getElem! _ _ α _ _ I xs n ← e | return .continue
let some n ← Nat.fromExpr? n | return .continue
let some xs ← getArrayLit? xs | return .continue
let r ← if h : n < xs.size then pure xs[n] else mkDefault α
return .done r

end Array
7 changes: 7 additions & 0 deletions tests/lean/run/array_simp.lean
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
#check_simp #[1,2,3,4,5][2] ~> 3
#check_simp #[1,2,3,4,5][2]? ~> some 3
#check_simp #[1,2,3,4,5][7]? ~> none
#check_simp #[][0]? ~> none
#check_simp #[1,2,3,4,5][2]! ~> 3
#check_simp #[1,2,3,4,5][7]! ~> (default : Nat)
#check_simp (#[] : Array Nat)[0]! ~> (default : Nat)

0 comments on commit af03af5

Please sign in to comment.