Make progress on fixing the extraction for Lean

5 files changed, 180 insertions, 220 deletions

diff --git a/backends/lean/Base/Primitives/ArraySlice.lean b/backends/lean/Base/Primitives/ArraySlice.lean
index 615e0783..2a080ca6 100644
--- a/backends/lean/Base/Primitives/ArraySlice.lean
+++ b/backends/lean/Base/Primitives/ArraySlice.lean
@@ -40,14 +40,14 @@ def Array.make (α : Type u) (n : Usize) (init : List α) (hl : init.len = n.val
 example : Array Int (Usize.ofInt 2) := Array.make Int (Usize.ofInt 2) [0, 1]
 
 @[simp]
-abbrev Array.index {α : Type u} {n : Usize} [Inhabited α] (v : Array α n) (i : Int) : α :=
+abbrev Array.index_s {α : Type u} {n : Usize} [Inhabited α] (v : Array α n) (i : Int) : α :=
   v.val.index i
 
 @[simp]
 abbrev Array.slice {α : Type u} {n : Usize} [Inhabited α] (v : Array α n) (i j : Int) : List α :=
   v.val.slice i j
 
-def Array.index_shared (α : Type u) (n : Usize) (v: Array α n) (i: Usize) : Result α :=
+def Array.index_usize (α : Type u) (n : Usize) (v: Array α n) (i: Usize) : Result α :=
   match v.val.indexOpt i.val with
   | none => fail .arrayOutOfBounds
   | some x => ret x
@@ -67,48 +67,27 @@ theorem Array.repeat_spec {α : Type u} (n : Usize) (x : α) :
  -/
 
 @[pspec]
-theorem Array.index_shared_spec {α : Type u} {n : Usize} [Inhabited α] (v: Array α n) (i: Usize)
+theorem Array.index_usize_spec {α : Type u} {n : Usize} [Inhabited α] (v: Array α n) (i: Usize)
   (hbound : i.val < v.length) :
-  ∃ x, v.index_shared α n i = ret x ∧ x = v.val.index i.val := by
-  simp only [index_shared]
+  ∃ x, v.index_usize α n i = ret x ∧ x = v.val.index i.val := by
+  simp only [index_usize]
   -- TODO: dependent rewrite
   have h := List.indexOpt_eq_index v.val i.val (by scalar_tac) (by simp [*])
   simp [*]
 
--- This shouldn't be used
-def Array.index_shared_back (α : Type u) (n : Usize) (v: Array α n) (i: Usize) (_: α) : Result Unit :=
-  if i.val < List.length v.val then
-    .ret ()
-  else
-    .fail arrayOutOfBounds
-
-def Array.index_mut (α : Type u) (n : Usize) (v: Array α n) (i: Usize) : Result α :=
-  match v.val.indexOpt i.val with
-  | none => fail .arrayOutOfBounds
-  | some x => ret x
-
-@[pspec]
-theorem Array.index_mut_spec {α : Type u} {n : Usize} [Inhabited α] (v: Array α n) (i: Usize)
-  (hbound : i.val < v.length) :
-  ∃ x, v.index_mut α n i = ret x ∧ x = v.val.index i.val := by
-  simp only [index_mut]
-  -- TODO: dependent rewrite
-  have h := List.indexOpt_eq_index v.val i.val (by scalar_tac) (by simp [*])
-  simp [*]
-
-def Array.index_mut_back (α : Type u) (n : Usize) (v: Array α n) (i: Usize) (x: α) : Result (Array α n) :=
+def Array.update_usize (α : Type u) (n : Usize) (v: Array α n) (i: Usize) (x: α) : Result (Array α n) :=
   match v.val.indexOpt i.val with
   | none => fail .arrayOutOfBounds
   | some _ =>
     .ret ⟨ v.val.update i.val x, by have := v.property; simp [*] ⟩
 
 @[pspec]
-theorem Array.index_mut_back_spec {α : Type u} {n : Usize} (v: Array α n) (i: Usize) (x : α)
+theorem Array.update_usize_spec {α : Type u} {n : Usize} (v: Array α n) (i: Usize) (x : α)
   (hbound : i.val < v.length) :
-  ∃ nv, v.index_mut_back α n i x = ret nv ∧
+  ∃ nv, v.update_usize α n i x = ret nv ∧
   nv.val = v.val.update i.val x
   := by
-  simp only [index_mut_back]
+  simp only [update_usize]
   have h := List.indexOpt_bounds v.val i.val
   split
   . simp_all [length]; cases h <;> scalar_tac
@@ -144,14 +123,14 @@ theorem Slice.len_val {α : Type u} (v : Slice α) : (Slice.len α v).val = v.le
   by rfl
 
 @[simp]
-abbrev Slice.index {α : Type u} [Inhabited α] (v: Slice α) (i: Int) : α :=
+abbrev Slice.index_s {α : Type u} [Inhabited α] (v: Slice α) (i: Int) : α :=
   v.val.index i
 
 @[simp]
 abbrev Slice.slice {α : Type u} [Inhabited α] (s : Slice α) (i j : Int) : List α :=
   s.val.slice i j
 
-def Slice.index_shared (α : Type u) (v: Slice α) (i: Usize) : Result α :=
+def Slice.index_usize (α : Type u) (v: Slice α) (i: Usize) : Result α :=
   match v.val.indexOpt i.val with
   | none => fail .arrayOutOfBounds
   | some x => ret x
@@ -162,10 +141,10 @@ def Slice.index_shared (α : Type u) (v: Slice α) (i: Usize) : Result α :=
  -/
 
 @[pspec]
-theorem Slice.index_shared_spec {α : Type u} [Inhabited α] (v: Slice α) (i: Usize)
+theorem Slice.index_usize_spec {α : Type u} [Inhabited α] (v: Slice α) (i: Usize)
   (hbound : i.val < v.length) :
-  ∃ x, v.index_shared α i = ret x ∧ x = v.val.index i.val := by
-  simp only [index_shared]
+  ∃ x, v.index_usize α i = ret x ∧ x = v.val.index i.val := by
+  simp only [index_usize]
   -- TODO: dependent rewrite
   have h := List.indexOpt_eq_index v.val i.val (by scalar_tac) (by simp [*])
   simp [*]
@@ -177,33 +156,19 @@ def Slice.index_shared_back (α : Type u) (v: Slice α) (i: Usize) (_: α) : Res
   else
     .fail arrayOutOfBounds
 
-def Slice.index_mut (α : Type u) (v: Slice α) (i: Usize) : Result α :=
-  match v.val.indexOpt i.val with
-  | none => fail .arrayOutOfBounds
-  | some x => ret x
-
-@[pspec]
-theorem Slice.index_mut_spec {α : Type u} [Inhabited α] (v: Slice α) (i: Usize)
-  (hbound : i.val < v.length) :
-  ∃ x, v.index_mut α i = ret x ∧ x = v.val.index i.val := by
-  simp only [index_mut]
-  -- TODO: dependent rewrite
-  have h := List.indexOpt_eq_index v.val i.val (by scalar_tac) (by simp [*])
-  simp [*]
-
-def Slice.index_mut_back (α : Type u) (v: Slice α) (i: Usize) (x: α) : Result (Slice α) :=
+def Slice.update_usize (α : Type u) (v: Slice α) (i: Usize) (x: α) : Result (Slice α) :=
   match v.val.indexOpt i.val with
   | none => fail .arrayOutOfBounds
   | some _ =>
     .ret ⟨ v.val.update i.val x, by have := v.property; simp [*] ⟩
 
 @[pspec]
-theorem Slice.index_mut_back_spec {α : Type u} (v: Slice α) (i: Usize) (x : α)
+theorem Slice.update_usize_spec {α : Type u} (v: Slice α) (i: Usize) (x : α)
   (hbound : i.val < v.length) :
-  ∃ nv, v.index_mut_back α i x = ret nv ∧
+  ∃ nv, v.update_usize α i x = ret nv ∧
   nv.val = v.val.update i.val x
   := by
-  simp only [index_mut_back]
+  simp only [update_usize]
   have h := List.indexOpt_bounds v.val i.val
   split
   . simp_all [length]; cases h <;> scalar_tac
@@ -212,34 +177,27 @@ theorem Slice.index_mut_back_spec {α : Type u} (v: Slice α) (i: Usize) (x : α
 /- Array to slice/subslices -/
 
 /- We could make this function not use the `Result` type. By making it monadic, we
-   push the user to use the `Array.to_slice_shared_spec` spec theorem below (through the
-   `progress` tactic), meaning `Array.to_slice_shared` should be considered as opaque.
+   push the user to use the `Array.to_slice_spec` spec theorem below (through the
+   `progress` tactic), meaning `Array.to_slice` should be considered as opaque.
    All what the spec theorem reveals is that the "representative" lists are the same. -/
-def Array.to_slice_shared (α : Type u) (n : Usize) (v : Array α n) : Result (Slice α) :=
+def Array.to_slice (α : Type u) (n : Usize) (v : Array α n) : Result (Slice α) :=
   ret ⟨ v.val, by simp [← List.len_eq_length]; scalar_tac ⟩
 
 @[pspec]
-theorem Array.to_slice_shared_spec {α : Type u} {n : Usize} (v : Array α n) :
-  ∃ s, to_slice_shared α n v = ret s ∧ v.val = s.val := by simp [to_slice_shared]
-
-def Array.to_slice_mut (α : Type u) (n : Usize) (v : Array α n) : Result (Slice α) :=
-  to_slice_shared α n v
+theorem Array.to_slice_spec {α : Type u} {n : Usize} (v : Array α n) :
+  ∃ s, to_slice α n v = ret s ∧ v.val = s.val := by simp [to_slice]
 
-@[pspec]
-theorem Array.to_slice_mut_spec {α : Type u} {n : Usize} (v : Array α n) :
-  ∃ s, Array.to_slice_shared α n v = ret s ∧ v.val = s.val := to_slice_shared_spec v
-
-def Array.to_slice_mut_back (α : Type u) (n : Usize) (_ : Array α n) (s : Slice α) : Result (Array α n) :=
+def Array.from_slice (α : Type u) (n : Usize) (_ : Array α n) (s : Slice α) : Result (Array α n) :=
   if h: s.val.len = n.val then
     ret ⟨ s.val, by simp [← List.len_eq_length, *] ⟩
   else fail panic
 
 @[pspec]
-theorem Array.to_slice_mut_back_spec {α : Type u} {n : Usize} (a : Array α n) (ns : Slice α) (h : ns.val.len = n.val) :
-  ∃ na, to_slice_mut_back α n a ns = ret na ∧ na.val = ns.val
-  := by simp [to_slice_mut_back, *]
+theorem Array.from_slice_spec {α : Type u} {n : Usize} (a : Array α n) (ns : Slice α) (h : ns.val.len = n.val) :
+  ∃ na, from_slice α n a ns = ret na ∧ na.val = ns.val
+  := by simp [from_slice, *]
 
-def Array.subslice_shared (α : Type u) (n : Usize) (a : Array α n) (r : Range Usize) : Result (Slice α) :=
+def Array.subslice (α : Type u) (n : Usize) (a : Array α n) (r : Range Usize) : Result (Slice α) :=
   -- TODO: not completely sure here
   if r.start.val < r.end_.val ∧ r.end_.val ≤ a.val.len then
     ret ⟨ a.val.slice r.start.val r.end_.val,
@@ -251,29 +209,18 @@ def Array.subslice_shared (α : Type u) (n : Usize) (a : Array α n) (r : Range
     fail panic
 
 @[pspec]
-theorem Array.subslice_shared_spec {α : Type u} {n : Usize} [Inhabited α] (a : Array α n) (r : Range Usize)
+theorem Array.subslice_spec {α : Type u} {n : Usize} [Inhabited α] (a : Array α n) (r : Range Usize)
   (h0 : r.start.val < r.end_.val) (h1 : r.end_.val ≤ a.val.len) :
-  ∃ s, subslice_shared α n a r = ret s ∧
+  ∃ s, subslice α n a r = ret s ∧
   s.val = a.val.slice r.start.val r.end_.val ∧
   (∀ i, 0 ≤ i → i + r.start.val < r.end_.val → s.val.index i = a.val.index (r.start.val + i))
   := by
-  simp [subslice_shared, *]
+  simp [subslice, *]
   intro i _ _
   have := List.index_slice r.start.val r.end_.val i a.val (by scalar_tac) (by scalar_tac) (by trivial) (by scalar_tac)
   simp [*]
 
-def Array.subslice_mut (α : Type u) (n : Usize) (a : Array α n) (r : Range Usize) : Result (Slice α) :=
-  Array.subslice_shared α n a r
-
-@[pspec]
-theorem Array.subslice_mut_spec {α : Type u} {n : Usize} [Inhabited α] (a : Array α n) (r : Range Usize)
-  (h0 : r.start.val < r.end_.val) (h1 : r.end_.val ≤ a.val.len) :
-  ∃ s, subslice_mut α n a r = ret s ∧
-  s.val = a.slice r.start.val r.end_.val ∧
-  (∀ i, 0 ≤ i → i + r.start.val < r.end_.val → s.val.index i = a.val.index (r.start.val + i))
-  := subslice_shared_spec a r h0 h1
-
-def Array.subslice_mut_back (α : Type u) (n : Usize) (a : Array α n) (r : Range Usize) (s : Slice α) : Result (Array α n) :=
+def Array.update_subslice (α : Type u) (n : Usize) (a : Array α n) (r : Range Usize) (s : Slice α) : Result (Array α n) :=
   -- TODO: not completely sure here
   if h: r.start.val < r.end_.val ∧ r.end_.val ≤ a.length ∧ s.val.len = r.end_.val - r.start.val then
     let s_beg := a.val.itake r.start.val
@@ -298,13 +245,13 @@ def Array.subslice_mut_back (α : Type u) (n : Usize) (a : Array α n) (r : Rang
 -- We should introduce special symbols for the monadic arithmetic operations
 -- (the use will never write those symbols directly).
 @[pspec]
-theorem Array.subslice_mut_back_spec {α : Type u} {n : Usize} [Inhabited α] (a : Array α n) (r : Range Usize) (s : Slice α)
+theorem Array.update_subslice_spec {α : Type u} {n : Usize} [Inhabited α] (a : Array α n) (r : Range Usize) (s : Slice α)
   (_ : r.start.val < r.end_.val) (_ : r.end_.val ≤ a.length) (_ : s.length = r.end_.val - r.start.val) :
-  ∃ na, subslice_mut_back α n a r s = ret na ∧
-  (∀ i, 0 ≤ i → i < r.start.val → na.index i = a.index i) ∧
-  (∀ i, r.start.val ≤ i → i < r.end_.val → na.index i = s.index (i - r.start.val)) ∧
-  (∀ i, r.end_.val ≤ i → i < n.val → na.index i = a.index i) := by
-  simp [subslice_mut_back, *]
+  ∃ na, update_subslice α n a r s = ret na ∧
+  (∀ i, 0 ≤ i → i < r.start.val → na.index_s i = a.index_s i) ∧
+  (∀ i, r.start.val ≤ i → i < r.end_.val → na.index_s i = s.index_s (i - r.start.val)) ∧
+  (∀ i, r.end_.val ≤ i → i < n.val → na.index_s i = a.index_s i) := by
+  simp [update_subslice, *]
   have h := List.replace_slice_index r.start.val r.end_.val a.val s.val
     (by scalar_tac) (by scalar_tac) (by scalar_tac) (by scalar_tac)
   simp [List.replace_slice] at h
@@ -321,7 +268,7 @@ theorem Array.subslice_mut_back_spec {α : Type u} {n : Usize} [Inhabited α] (a
     have := h2 i (by int_tac) (by int_tac)
     simp [*]
 
-def Slice.subslice_shared (α : Type u) (s : Slice α) (r : Range Usize) : Result (Slice α) :=
+def Slice.subslice (α : Type u) (s : Slice α) (r : Range Usize) : Result (Slice α) :=
   -- TODO: not completely sure here
   if r.start.val < r.end_.val ∧ r.end_.val ≤ s.length then
     ret ⟨ s.val.slice r.start.val r.end_.val,
@@ -333,32 +280,21 @@ def Slice.subslice_shared (α : Type u) (s : Slice α) (r : Range Usize) : Resul
     fail panic
 
 @[pspec]
-theorem Slice.subslice_shared_spec {α : Type u} [Inhabited α] (s : Slice α) (r : Range Usize)
+theorem Slice.subslice_spec {α : Type u} [Inhabited α] (s : Slice α) (r : Range Usize)
   (h0 : r.start.val < r.end_.val) (h1 : r.end_.val ≤ s.val.len) :
-  ∃ ns, subslice_shared α s r = ret ns ∧
+  ∃ ns, subslice α s r = ret ns ∧
   ns.val = s.slice r.start.val r.end_.val ∧
-  (∀ i, 0 ≤ i → i + r.start.val < r.end_.val → ns.index i = s.index (r.start.val + i))
+  (∀ i, 0 ≤ i → i + r.start.val < r.end_.val → ns.index_s i = s.index_s (r.start.val + i))
   := by
-  simp [subslice_shared, *]
+  simp [subslice, *]
   intro i _ _
   have := List.index_slice r.start.val r.end_.val i s.val (by scalar_tac) (by scalar_tac) (by trivial) (by scalar_tac)
   simp [*]
 
-def Slice.subslice_mut (α : Type u) (s : Slice α) (r : Range Usize) : Result (Slice α) :=
-  Slice.subslice_shared α s r
-
-@[pspec]
-theorem Slice.subslice_mut_spec {α : Type u} [Inhabited α] (s : Slice α) (r : Range Usize)
-  (h0 : r.start.val < r.end_.val) (h1 : r.end_.val ≤ s.val.len) :
-  ∃ ns, subslice_mut α s r = ret ns ∧
-  ns.val = s.slice r.start.val r.end_.val ∧
-  (∀ i, 0 ≤ i → i + r.start.val < r.end_.val → ns.index i = s.index (r.start.val + i))
-  := subslice_shared_spec s r h0 h1
-
 attribute [pp_dot] List.len List.length List.index -- use the dot notation when printing
 set_option pp.coercions false -- do not print coercions with ↑ (this doesn't parse)
 
-def Slice.subslice_mut_back (α : Type u) (s : Slice α) (r : Range Usize) (ss : Slice α) : Result (Slice α) :=
+def Slice.update_subslice (α : Type u) (s : Slice α) (r : Range Usize) (ss : Slice α) : Result (Slice α) :=
   -- TODO: not completely sure here
   if h: r.start.val < r.end_.val ∧ r.end_.val ≤ s.length ∧ ss.val.len = r.end_.val - r.start.val then
     let s_beg := s.val.itake r.start.val
@@ -378,13 +314,13 @@ def Slice.subslice_mut_back (α : Type u) (s : Slice α) (r : Range Usize) (ss :
     fail panic
 
 @[pspec]
-theorem Slice.subslice_mut_back_spec {α : Type u} [Inhabited α] (a : Slice α) (r : Range Usize) (ss : Slice α)
+theorem Slice.update_subslice_spec {α : Type u} [Inhabited α] (a : Slice α) (r : Range Usize) (ss : Slice α)
   (_ : r.start.val < r.end_.val) (_ : r.end_.val ≤ a.length) (_ : ss.length = r.end_.val - r.start.val) :
-  ∃ na, subslice_mut_back α a r ss = ret na ∧
-  (∀ i, 0 ≤ i → i < r.start.val → na.index i = a.index i) ∧
-  (∀ i, r.start.val ≤ i → i < r.end_.val → na.index i = ss.index (i - r.start.val)) ∧
-  (∀ i, r.end_.val ≤ i → i < a.length → na.index i = a.index i) := by
-  simp [subslice_mut_back, *]
+  ∃ na, update_subslice α a r ss = ret na ∧
+  (∀ i, 0 ≤ i → i < r.start.val → na.index_s i = a.index_s i) ∧
+  (∀ i, r.start.val ≤ i → i < r.end_.val → na.index_s i = ss.index_s (i - r.start.val)) ∧
+  (∀ i, r.end_.val ≤ i → i < a.length → na.index_s i = a.index_s i) := by
+  simp [update_subslice, *]
   have h := List.replace_slice_index r.start.val r.end_.val a.val ss.val
     (by scalar_tac) (by scalar_tac) (by scalar_tac) (by scalar_tac)
   simp [List.replace_slice, *] at h
diff --git a/compiler/Extract.ml b/compiler/Extract.ml
index fdcd82d9..574602c7 100644
--- a/compiler/Extract.ml
+++ b/compiler/Extract.ml
@@ -241,7 +241,7 @@ let rec extract_typed_pattern (ctx : extraction_ctx) (fmt : F.formatter)
       ctx
   | PatVar (v, _) ->
       let vname =
-        ctx.fmt.var_basename ctx.names_map.names_set v.basename v.ty
+        ctx.fmt.var_basename ctx.names_maps.names_map.names_set v.basename v.ty
       in
       let ctx, vname = ctx_add_var vname v.id ctx in
       F.pp_print_string fmt vname;
diff --git a/compiler/ExtractBase.ml b/compiler/ExtractBase.ml
index 7e8e4ffc..8f71116c 100644
--- a/compiler/ExtractBase.ml
+++ b/compiler/ExtractBase.ml
@@ -520,23 +520,6 @@ let names_map_add (id_to_string : id -> string) (id : id) (name : string)
   let names_set = StringSet.add name nm.names_set in
   { id_to_name; name_to_id; names_set }
 
-let names_map_add_assumed_type (id_to_string : id -> string) (id : assumed_ty)
-    (name : string) (nm : names_map) : names_map =
-  names_map_add id_to_string (TypeId (Assumed id)) name nm
-
-let names_map_add_assumed_struct (id_to_string : id -> string) (id : assumed_ty)
-    (name : string) (nm : names_map) : names_map =
-  names_map_add id_to_string (StructId (Assumed id)) name nm
-
-let names_map_add_assumed_variant (id_to_string : id -> string)
-    (id : assumed_ty) (variant_id : VariantId.id) (name : string)
-    (nm : names_map) : names_map =
-  names_map_add id_to_string (VariantId (Assumed id, variant_id)) name nm
-
-let names_map_add_function (id_to_string : id -> string) (fid : fun_id)
-    (name : string) (nm : names_map) : names_map =
-  names_map_add id_to_string (FunId fid) name nm
-
 (** The unsafe names map stores mappings from identifiers to names which might
     collide. For some backends and some names, it might be acceptable to have
     collisions. For instance, in Lean, different records can have fields with
@@ -547,6 +530,8 @@ let names_map_add_function (id_to_string : id -> string) (fid : fun_id)
   *)
 type unsafe_names_map = { id_to_name : string IdMap.t }
 
+let empty_unsafe_names_map = { id_to_name = IdMap.empty }
+
 let unsafe_names_map_add (id : id) (name : string) (nm : unsafe_names_map) :
     unsafe_names_map =
   { id_to_name = IdMap.add id name nm.id_to_name }
@@ -585,16 +570,7 @@ let basename_to_unique (names_set : StringSet.t)
 
 type fun_name_info = { keep_fwd : bool; num_backs : int }
 
-(** Extraction context.
-
-    Note that the extraction context contains information coming from the
-    LLBC AST (not only the pure AST). This is useful for naming, for instance:
-    we use the region information to generate the names of the backward
-    functions, etc.
- *)
-type extraction_ctx = {
-  crate : A.crate;
-  trans_ctx : trans_ctx;
+type names_maps = {
   names_map : names_map;
       (** The map for id to names, where we forbid name collisions
           (ex.: we always forbid function name collisions). *)
@@ -610,6 +586,19 @@ type extraction_ctx = {
           the name "u32", and another field of the same record refers to "u32"
           (for instance in its type).
        *)
+}
+
+(** Extraction context.
+
+    Note that the extraction context contains information coming from the
+    LLBC AST (not only the pure AST). This is useful for naming, for instance:
+    we use the region information to generate the names of the backward
+    functions, etc.
+ *)
+type extraction_ctx = {
+  crate : A.crate;
+  trans_ctx : trans_ctx;
+  names_maps : names_maps;
   fmt : formatter;
   indent_incr : int;
       (** The indent increment we insert whenever we need to indent more *)
@@ -836,12 +825,15 @@ let allow_collisions (id : id) : bool =
   | FieldId _ | TraitItemClauseId _ | TraitParentClauseId _ | TraitItemId _
   | TraitMethodId _ ->
       !Config.record_fields_short_names
+  | FunId (Pure _ | FromLlbc (FunId (Assumed _), _, _)) ->
+      (* We map several assumed functions to the same id *)
+      true
   | _ -> false
 
-let ctx_add (id : id) (name : string) (ctx : extraction_ctx) : extraction_ctx =
-  (* The id_to_string function to print nice debugging messages if there are
-   * collisions *)
-  let id_to_string (id : id) : string = id_to_string id ctx in
+(** The [id_to_string] function to print nice debugging messages if there are
+    collisions *)
+let names_maps_add (id_to_string : id -> string) (id : id) (name : string)
+    (nm : names_maps) : names_maps =
   (* We do not use the same name map if we allow/disallow collisions.
      We notably use it for field names: some backends like Lean can use the
      type information to disambiguate field projections.
@@ -856,59 +848,90 @@ let ctx_add (id : id) (name : string) (ctx : extraction_ctx) : extraction_ctx =
   *)
   if allow_collisions id then (
     (* Check with the ids which are considered to be strict on collisions *)
-    names_map_check_collision id_to_string id name ctx.strict_names_map;
+    names_map_check_collision id_to_string id name nm.strict_names_map;
     {
-      ctx with
-      unsafe_names_map = unsafe_names_map_add id name ctx.unsafe_names_map;
+      nm with
+      unsafe_names_map = unsafe_names_map_add id name nm.unsafe_names_map;
     })
   else
     (* Remark: if we are strict on collisions:
        - we add the id to the strict collisions map
        - we check that the id doesn't collide with the unsafe map
+       TODO: we might not check that:
+       - a user defined function doesn't collide with an assumed function
+       - two trait decl items don't collide with each other
     *)
     let strict_names_map =
       if strict_collisions id then
-        names_map_add id_to_string id name ctx.strict_names_map
-      else ctx.strict_names_map
+        names_map_add id_to_string id name nm.strict_names_map
+      else nm.strict_names_map
     in
-    let names_map = names_map_add id_to_string id name ctx.names_map in
-    { ctx with strict_names_map; names_map }
+    let names_map = names_map_add id_to_string id name nm.names_map in
+    { nm with strict_names_map; names_map }
 
-let ctx_get (id : id) (ctx : extraction_ctx) : string =
+let ctx_add (id : id) (name : string) (ctx : extraction_ctx) : extraction_ctx =
+  let id_to_string (id : id) : string = id_to_string id ctx in
+  let names_maps = names_maps_add id_to_string id name ctx.names_maps in
+  { ctx with names_maps }
+
+(** The [id_to_string] function to print nice debugging messages if there are
+    collisions *)
+let names_maps_get (id_to_string : id -> string) (id : id) (nm : names_maps) :
+    string =
   (* We do not use the same name map if we allow/disallow collisions *)
   let map_to_string (m : string IdMap.t) : string =
     "[\n"
     ^ String.concat ","
         (List.map
-           (fun (id, n) -> "\n  " ^ id_to_string id ctx ^ " -> " ^ n)
+           (fun (id, n) -> "\n  " ^ id_to_string id ^ " -> " ^ n)
            (IdMap.bindings m))
     ^ "\n]"
   in
   if allow_collisions id then (
-    let m = ctx.unsafe_names_map.id_to_name in
+    let m = nm.unsafe_names_map.id_to_name in
     match IdMap.find_opt id m with
     | Some s -> s
     | None ->
         let err =
-          "Could not find: " ^ id_to_string id ctx ^ "\nNames map:\n"
+          "Could not find: " ^ id_to_string id ^ "\nNames map:\n"
           ^ map_to_string m
         in
         log#serror err;
         if !Config.extract_fail_hard then raise (Failure err)
-        else
-          "(%%%ERROR: unknown identifier\": " ^ id_to_string id ctx ^ "\"%%%)")
+        else "(%%%ERROR: unknown identifier\": " ^ id_to_string id ^ "\"%%%)")
   else
-    let m = ctx.names_map.id_to_name in
+    let m = nm.names_map.id_to_name in
     match IdMap.find_opt id m with
     | Some s -> s
     | None ->
         let err =
-          "Could not find: " ^ id_to_string id ctx ^ "\nNames map:\n"
+          "Could not find: " ^ id_to_string id ^ "\nNames map:\n"
           ^ map_to_string m
         in
         log#serror err;
         if !Config.extract_fail_hard then raise (Failure err)
-        else "(ERROR: \"" ^ id_to_string id ctx ^ "\")"
+        else "(ERROR: \"" ^ id_to_string id ^ "\")"
+
+let ctx_get (id : id) (ctx : extraction_ctx) : string =
+  let id_to_string (id : id) : string = id_to_string id ctx in
+  names_maps_get id_to_string id ctx.names_maps
+
+let names_maps_add_assumed_type (id_to_string : id -> string) (id : assumed_ty)
+    (name : string) (nm : names_maps) : names_maps =
+  names_maps_add id_to_string (TypeId (Assumed id)) name nm
+
+let names_maps_add_assumed_struct (id_to_string : id -> string)
+    (id : assumed_ty) (name : string) (nm : names_maps) : names_maps =
+  names_maps_add id_to_string (StructId (Assumed id)) name nm
+
+let names_maps_add_assumed_variant (id_to_string : id -> string)
+    (id : assumed_ty) (variant_id : VariantId.id) (name : string)
+    (nm : names_maps) : names_maps =
+  names_maps_add id_to_string (VariantId (Assumed id, variant_id)) name nm
+
+let names_maps_add_function (id_to_string : id -> string) (fid : fun_id)
+    (name : string) (nm : names_maps) : names_maps =
+  names_maps_add id_to_string (FunId fid) name nm
 
 let ctx_get_global (id : A.GlobalDeclId.id) (ctx : extraction_ctx) : string =
   ctx_get (GlobalId id) ctx
@@ -999,9 +1022,12 @@ let ctx_get_termination_measure (def_id : A.FunDeclId.id)
 (** Generate a unique type variable name and add it to the context *)
 let ctx_add_type_var (basename : string) (id : TypeVarId.id)
     (ctx : extraction_ctx) : extraction_ctx * string =
-  let name = ctx.fmt.type_var_basename ctx.names_map.names_set basename in
   let name =
-    basename_to_unique ctx.names_map.names_set ctx.fmt.append_index name
+    ctx.fmt.type_var_basename ctx.names_maps.names_map.names_set basename
+  in
+  let name =
+    basename_to_unique ctx.names_maps.names_map.names_set ctx.fmt.append_index
+      name
   in
   let ctx = ctx_add (TypeVarId id) name ctx in
   (ctx, name)
@@ -1010,10 +1036,12 @@ let ctx_add_type_var (basename : string) (id : TypeVarId.id)
 let ctx_add_const_generic_var (basename : string) (id : ConstGenericVarId.id)
     (ctx : extraction_ctx) : extraction_ctx * string =
   let name =
-    ctx.fmt.const_generic_var_basename ctx.names_map.names_set basename
+    ctx.fmt.const_generic_var_basename ctx.names_maps.names_map.names_set
+      basename
   in
   let name =
-    basename_to_unique ctx.names_map.names_set ctx.fmt.append_index name
+    basename_to_unique ctx.names_maps.names_map.names_set ctx.fmt.append_index
+      name
   in
   let ctx = ctx_add (ConstGenericVarId id) name ctx in
   (ctx, name)
@@ -1029,7 +1057,8 @@ let ctx_add_type_vars (vars : (string * TypeVarId.id) list)
 let ctx_add_var (basename : string) (id : VarId.id) (ctx : extraction_ctx) :
     extraction_ctx * string =
   let name =
-    basename_to_unique ctx.names_map.names_set ctx.fmt.append_index basename
+    basename_to_unique ctx.names_maps.names_map.names_set ctx.fmt.append_index
+      basename
   in
   let ctx = ctx_add (VarId id) name ctx in
   (ctx, name)
@@ -1038,7 +1067,8 @@ let ctx_add_var (basename : string) (id : VarId.id) (ctx : extraction_ctx) :
 let ctx_add_trait_self_clause (ctx : extraction_ctx) : extraction_ctx * string =
   let basename = ctx.fmt.trait_self_clause_basename in
   let name =
-    basename_to_unique ctx.names_map.names_set ctx.fmt.append_index basename
+    basename_to_unique ctx.names_maps.names_map.names_set ctx.fmt.append_index
+      basename
   in
   let ctx = ctx_add TraitSelfClauseId name ctx in
   (ctx, name)
@@ -1047,7 +1077,8 @@ let ctx_add_trait_self_clause (ctx : extraction_ctx) : extraction_ctx * string =
 let ctx_add_local_trait_clause (basename : string) (id : TraitClauseId.id)
     (ctx : extraction_ctx) : extraction_ctx * string =
   let name =
-    basename_to_unique ctx.names_map.names_set ctx.fmt.append_index basename
+    basename_to_unique ctx.names_maps.names_map.names_set ctx.fmt.append_index
+      basename
   in
   let ctx = ctx_add (LocalTraitClauseId id) name ctx in
   (ctx, name)
@@ -1057,7 +1088,9 @@ let ctx_add_vars (vars : var list) (ctx : extraction_ctx) :
     extraction_ctx * string list =
   List.fold_left_map
     (fun ctx (v : var) ->
-      let name = ctx.fmt.var_basename ctx.names_map.names_set v.basename v.ty in
+      let name =
+        ctx.fmt.var_basename ctx.names_maps.names_map.names_set v.basename v.ty
+      in
       ctx_add_var name v.id ctx)
     ctx vars
 
@@ -1078,7 +1111,9 @@ let ctx_add_local_trait_clauses (clauses : trait_clause list)
     (ctx : extraction_ctx) : extraction_ctx * string list =
   List.fold_left_map
     (fun ctx (c : trait_clause) ->
-      let basename = ctx.fmt.trait_clause_basename ctx.names_map.names_set c in
+      let basename =
+        ctx.fmt.trait_clause_basename ctx.names_maps.names_map.names_set c
+      in
       ctx_add_local_trait_clause basename c.clause_id ctx)
     ctx clauses
 
@@ -1189,9 +1224,10 @@ type names_map_init = {
   assumed_pure_functions : (pure_assumed_fun_id * string) list;
 }
 
-(** Initialize a names map with a proper set of keywords/names coming from the
+(** Initialize names maps with a proper set of keywords/names coming from the
     target language/prover. *)
-let initialize_names_map (fmt : formatter) (init : names_map_init) : names_map =
+let initialize_names_maps (fmt : formatter) (init : names_map_init) : names_maps
+    =
   let int_names = List.map fmt.int_name T.all_int_types in
   let keywords =
     List.concat
@@ -1207,7 +1243,10 @@ let initialize_names_map (fmt : formatter) (init : names_map_init) : names_map =
    * Also note that we don't need this mapping for keywords: we insert keywords only
    * to check collisions. *)
   let id_to_name = IdMap.empty in
-  let nm = { id_to_name; name_to_id; names_set } in
+  let names_map = { id_to_name; name_to_id; names_set } in
+  let unsafe_names_map = empty_unsafe_names_map in
+  let strict_names_map = empty_names_map in
+  let nm = { names_map; unsafe_names_map; strict_names_map } in
   (* For debugging - we are creating bindings for assumed types and functions, so
    * it is ok if we simply use the "show" function (those aren't simply identified
    * by numbers) *)
@@ -1221,19 +1260,19 @@ let initialize_names_map (fmt : formatter) (init : names_map_init) : names_map =
   let nm =
     List.fold_left
       (fun nm (type_id, name) ->
-        names_map_add_assumed_type id_to_string type_id name nm)
+        names_maps_add_assumed_type id_to_string type_id name nm)
       nm init.assumed_adts
   in
   let nm =
     List.fold_left
       (fun nm (type_id, name) ->
-        names_map_add_assumed_struct id_to_string type_id name nm)
+        names_maps_add_assumed_struct id_to_string type_id name nm)
       nm init.assumed_structs
   in
   let nm =
     List.fold_left
       (fun nm (type_id, variant_id, name) ->
-        names_map_add_assumed_variant id_to_string type_id variant_id name nm)
+        names_maps_add_assumed_variant id_to_string type_id variant_id name nm)
       nm init.assumed_variants
   in
   let assumed_functions =
@@ -1245,7 +1284,7 @@ let initialize_names_map (fmt : formatter) (init : names_map_init) : names_map =
   in
   let nm =
     List.fold_left
-      (fun nm (fid, name) -> names_map_add_function id_to_string fid name nm)
+      (fun nm (fid, name) -> names_maps_add_function id_to_string fid name nm)
       nm assumed_functions
   in
   (* Return *)
diff --git a/compiler/ExtractTypes.ml b/compiler/ExtractTypes.ml
index 219f273f..fd3baf0d 100644
--- a/compiler/ExtractTypes.ml
+++ b/compiler/ExtractTypes.ml
@@ -300,30 +300,30 @@ let assumed_llbc_functions () :
   match !backend with
   | FStar | Coq | HOL4 ->
       [
-        (ArrayIndexShared, None, "array_index_shared");
-        (ArrayIndexMut, None, "array_index_mut_fwd");
-        (ArrayIndexMut, rg0, "array_index_mut_back");
-        (ArrayToSliceShared, None, "array_to_slice_shared");
-        (ArrayToSliceMut, None, "array_to_slice_mut_fwd");
-        (ArrayToSliceMut, rg0, "array_to_slice_mut_back");
+        (ArrayIndexShared, None, "array_index_usize");
+        (ArrayIndexMut, None, "array_index_usize");
+        (ArrayIndexMut, rg0, "array_update_usize");
+        (ArrayToSliceShared, None, "array_to_slice");
+        (ArrayToSliceMut, None, "array_to_slice");
+        (ArrayToSliceMut, rg0, "array_from_slice");
         (ArrayRepeat, None, "array_repeat");
-        (SliceIndexShared, None, "slice_index_shared");
-        (SliceIndexMut, None, "slice_index_mut_fwd");
-        (SliceIndexMut, rg0, "slice_index_mut_back");
+        (SliceIndexShared, None, "slice_index_usize");
+        (SliceIndexMut, None, "slice_index_usize");
+        (SliceIndexMut, rg0, "slice_update_usize");
         (SliceLen, None, "slice_len");
       ]
   | Lean ->
       [
-        (ArrayIndexShared, None, "Array.index_shared");
-        (ArrayIndexMut, None, "Array.index_mut");
-        (ArrayIndexMut, rg0, "Array.index_mut_back");
-        (ArrayToSliceShared, None, "Array.to_slice_shared");
-        (ArrayToSliceMut, None, "Array.to_slice_mut");
-        (ArrayToSliceMut, rg0, "Array.to_slice_mut_back");
+        (ArrayIndexShared, None, "Array.index_usize");
+        (ArrayIndexMut, None, "Array.index_usize");
+        (ArrayIndexMut, rg0, "Array.update_usize");
+        (ArrayToSliceShared, None, "Array.to_slice");
+        (ArrayToSliceMut, None, "Array.to_slice");
+        (ArrayToSliceMut, rg0, "Array.from_slice");
         (ArrayRepeat, None, "Array.repeat");
-        (SliceIndexShared, None, "Slice.index_shared");
-        (SliceIndexMut, None, "Slice.index_mut");
-        (SliceIndexMut, rg0, "Slice.index_mut_back");
+        (SliceIndexShared, None, "Slice.index_usize");
+        (SliceIndexMut, None, "Slice.index_usize");
+        (SliceIndexMut, rg0, "Slice.update_usize");
         (SliceLen, None, "Slice.len");
       ]
 
@@ -941,11 +941,11 @@ let mk_formatter (ctx : trans_ctx) (crate_name : string)
     extract_binop;
   }
 
-let mk_formatter_and_names_map (ctx : trans_ctx) (crate_name : string)
-    (variant_concatenate_type_name : bool) : formatter * names_map =
+let mk_formatter_and_names_maps (ctx : trans_ctx) (crate_name : string)
+    (variant_concatenate_type_name : bool) : formatter * names_maps =
   let fmt = mk_formatter ctx crate_name variant_concatenate_type_name in
-  let names_map = initialize_names_map fmt (names_map_init ()) in
-  (fmt, names_map)
+  let names_maps = initialize_names_maps fmt (names_map_init ()) in
+  (fmt, names_maps)
 
 let is_single_opaque_fun_decl_group (dg : Pure.fun_decl list) : bool =
   match dg with [ d ] -> d.body = None | _ -> false
@@ -1507,8 +1507,8 @@ let extract_type_decl_variant (ctx : extraction_ctx) (fmt : F.formatter)
           | Some field_name ->
               let var_id = VarId.of_int (FieldId.to_int fid) in
               let field_name =
-                ctx.fmt.var_basename ctx.names_map.names_set (Some field_name)
-                  f.field_ty
+                ctx.fmt.var_basename ctx.names_maps.names_map.names_set
+                  (Some field_name) f.field_ty
               in
               let ctx, field_name = ctx_add_var field_name var_id ctx in
               F.pp_print_string fmt (field_name ^ " :");
diff --git a/compiler/Translate.ml b/compiler/Translate.ml
index c5ac4e96..cb23198a 100644
--- a/compiler/Translate.ml
+++ b/compiler/Translate.ml
@@ -990,23 +990,10 @@ let translate_crate (filename : string) (dest_dir : string) (crate : A.crate) :
   in
   (* Initialize the names map (we insert the names of the "primitives"
      declarations, and insert the names of the local declarations later) *)
-  let mk_formatter_and_names_map = Extract.mk_formatter_and_names_map in
-  let fmt, names_map =
-    mk_formatter_and_names_map trans_ctx crate.name
+  let fmt, names_maps =
+    Extract.mk_formatter_and_names_maps trans_ctx crate.name
       variant_concatenate_type_name
   in
-  let strict_names_map =
-    let open ExtractBase in
-    let ids =
-      List.filter
-        (fun (id, _) -> strict_collisions id)
-        (IdMap.bindings names_map.id_to_name)
-    in
-    List.fold_left
-      (* id_to_string: we shouldn't need to use it *)
-        (fun m (id, n) -> names_map_add show_id id n m)
-      empty_names_map ids
-  in
 
   (* We need to compute which functions are recursive, in order to know
    * whether we should generate a decrease clause or not. *)
@@ -1060,9 +1047,7 @@ let translate_crate (filename : string) (dest_dir : string) (crate : A.crate) :
     {
       ExtractBase.crate;
       trans_ctx;
-      names_map;
-      unsafe_names_map = { id_to_name = ExtractBase.IdMap.empty };
-      strict_names_map;
+      names_maps;
       fmt;
       indent_incr = 2;
       use_dep_ite = !Config.backend = Lean && !Config.extract_decreases_clauses;