feat: simprocs for #[1,2,3,4,5][2]

src/Lean/Meta/AppBuilder.lean

@@ -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)

src/Lean/Meta/LitValues.lean

@@ -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

src/Lean/Meta/Tactic/Simp/BuiltinSimprocs.lean

@@ -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

src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/Array.lean

@@ -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

tests/lean/run/array_simp.lean

@@ -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)