Fix minor issues

3 files changed, 83 insertions, 1 deletions

diff --git a/backends/coq/Primitives.v b/backends/coq/Primitives.v
index e6d3118f..990e27e4 100644
--- a/backends/coq/Primitives.v
+++ b/backends/coq/Primitives.v
@@ -585,6 +585,13 @@ Axiom array_repeat : forall (T : Type) (n : usize) (x : T), array T n.
 Axiom array_index_usize : forall (T : Type) (n : usize) (x : array T n) (i : usize), result T.
 Axiom array_update_usize : forall (T : Type) (n : usize) (x : array T n) (i : usize) (nx : T), result (array T n).
 
+Definition array_index_mut_usize (T : Type) (n : usize) (a : array T n) (i : usize) :
+  result (T * (T -> result (array T n))) :=
+  match array_index_usize T n a i with
+  | Fail_ e => Fail_ e
+  | Return x => Return (x, array_update_usize T n a i)
+  end.
+
 (*** Slice *)
 Definition slice T := { l: list T | Z.of_nat (length l) <= usize_max}.
 
@@ -592,11 +599,25 @@ Axiom slice_len : forall (T : Type) (s : slice T), usize.
 Axiom slice_index_usize : forall (T : Type) (x : slice T) (i : usize), result T.
 Axiom slice_update_usize : forall (T : Type) (x : slice T) (i : usize) (nx : T), result (slice T).
 
+Definition slice_index_mut_usize (T : Type) (s : slice T) (i : usize) :
+  result (T * (T -> result (slice T))) :=
+  match slice_index_usize T s i with
+  | Fail_ e => Fail_ e
+  | Return x => Return (x, slice_update_usize T s i)
+  end.
+
 (*** Subslices *)
 
 Axiom array_to_slice : forall (T : Type) (n : usize) (x : array T n), result (slice T).
 Axiom array_from_slice : forall (T : Type) (n : usize) (x : array T n) (s : slice T), result (array T n).
 
+Definition array_to_slice_mut (T : Type) (n : usize) (a : array T n) :
+  result (slice T * (slice T -> result (array T n))) :=
+  match array_to_slice T n a with
+  | Fail_ e => Fail_ e
+  | Return x => Return (x, array_from_slice T n a)
+  end.
+
 Axiom array_subslice: forall (T : Type) (n : usize) (x : array T n) (r : core_ops_range_Range usize), result (slice T).
 Axiom array_update_subslice: forall (T : Type) (n : usize) (x : array T n) (r : core_ops_range_Range usize) (ns : slice T), result (array T n).
 
diff --git a/backends/fstar/merge/Primitives.fst b/backends/fstar/merge/Primitives.fst
index 8011efa1..fca80829 100644
--- a/backends/fstar/merge/Primitives.fst
+++ b/backends/fstar/merge/Primitives.fst
@@ -531,10 +531,18 @@ let array_index_usize (a : Type0) (n : usize) (x : array a n) (i : usize) : resu
   if i < length x then Return (index x i)
   else Fail Failure
 
-let array_update_usize (a : Type0) (n : usize) (x : array a n) (i : usize) (nx : a) : result (array a n) =
+let array_update_usize (a : Type0) (n : usize) (x : array a n) (i : usize) (nx : a) :
+  result (array a n) =
   if i < length x then Return (list_update x i nx)
   else Fail Failure
 
+let array_index_mut_usize (a : Type0) (n : usize) (x : array a n) (i : usize) :
+  result (a & (a -> result (array a n))) =
+  match array_index_usize a n x i with
+  | Fail e -> Fail e
+  | Return v ->
+    Return (v, array_update_usize a n x i)
+
 (*** Slice *)
 type slice (a : Type0) = s:list a{length s <= usize_max}
 
@@ -548,6 +556,13 @@ let slice_update_usize (a : Type0) (x : slice a) (i : usize) (nx : a) : result (
   if i < length x then Return (list_update x i nx)
   else Fail Failure
 
+let slice_index_mut_usize (a : Type0) (s : slice a) (i : usize) :
+  result (a & (a -> result (slice a))) =
+  match slice_index_usize a s i with
+  | Fail e -> Fail e
+  | Return x ->
+    Return (x, slice_update_usize a s i)
+
 (*** Subslices *)
 
 let array_to_slice (a : Type0) (n : usize) (x : array a n) : result (slice a) = Return x
@@ -555,6 +570,10 @@ let array_from_slice (a : Type0) (n : usize) (x : array a n) (s : slice a) : res
   if length s = n then Return s
   else Fail Failure
 
+let array_to_slice_mut (a : Type0) (n : usize) (x : array a n) :
+  result (slice a & (slice a -> result (array a n))) =
+  Return (x, array_from_slice a n x)
+
 // TODO: finish the definitions below (there lacks [List.drop] and [List.take] in the standard library *)
 let array_subslice (a : Type0) (n : usize) (x : array a n) (r : core_ops_range_Range usize) : result (slice a) =
   admit()
diff --git a/backends/lean/Base/Primitives/ArraySlice.lean b/backends/lean/Base/Primitives/ArraySlice.lean
index 59432a0b..5057fb01 100644
--- a/backends/lean/Base/Primitives/ArraySlice.lean
+++ b/backends/lean/Base/Primitives/ArraySlice.lean
@@ -93,6 +93,21 @@ theorem Array.update_usize_spec {α : Type u} {n : Usize} (v: Array α n) (i: Us
   . simp_all [length]; cases h <;> scalar_tac
   . simp_all
 
+def Array.index_mut_usize (α : Type u) (n : Usize) (v: Array α n) (i: Usize) :
+  Result (α × (α -> Result (Array α n))) := do
+  let x ← index_usize α n v i
+  ret (x, update_usize α n v i)
+
+@[pspec]
+theorem Array.index_mut_usize_spec {α : Type u} {n : Usize} [Inhabited α] (v: Array α n) (i: Usize)
+  (hbound : i.val < v.length) :
+  ∃ x back, v.index_mut_usize α n i = ret (x, back) ∧
+  x = v.val.index i.val ∧
+  back = update_usize α n v i := by
+  simp only [index_mut_usize, Bind.bind, bind]
+  have ⟨ x, h ⟩ := index_usize_spec v i hbound
+  simp [h]
+
 def Slice (α : Type u) := { l : List α // l.length ≤ Usize.max }
 
 instance (a : Type u) : Arith.HasIntProp (Slice a) where
@@ -167,6 +182,21 @@ theorem Slice.update_usize_spec {α : Type u} (v: Slice α) (i: Usize) (x : α)
   . simp_all [length]; cases h <;> scalar_tac
   . simp_all
 
+def Slice.index_mut_usize (α : Type u) (v: Slice α) (i: Usize) :
+  Result (α × (α → Result (Slice α))) := do
+  let x ← Slice.index_usize α v i
+  ret (x, Slice.update_usize α v i)
+
+@[pspec]
+theorem Slice.index_mut_usize_spec {α : Type u} [Inhabited α] (v: Slice α) (i: Usize)
+  (hbound : i.val < v.length) :
+  ∃ x back, v.index_mut_usize α i = ret (x, back) ∧
+  x = v.val.index i.val ∧
+  back = Slice.update_usize α v i := by
+  simp only [index_mut_usize, Bind.bind, bind]
+  have ⟨ x, h ⟩ := Slice.index_usize_spec v i hbound
+  simp [h]
+
 /- Array to slice/subslices -/
 
 /- We could make this function not use the `Result` type. By making it monadic, we
@@ -190,6 +220,18 @@ theorem Array.from_slice_spec {α : Type u} {n : Usize} (a : Array α n) (ns : S
   ∃ na, from_slice α n a ns = ret na ∧ na.val = ns.val
   := by simp [from_slice, *]
 
+def Array.to_slice_mut (α : Type u) (n : Usize) (a : Array α n) :
+  Result (Slice α × (Slice α → Result (Array α n))) := do
+  let s ← Array.to_slice α n a
+  ret (s, Array.from_slice α n a)
+
+@[pspec]
+theorem Array.to_slice_mut_spec {α : Type u} {n : Usize} (v : Array α n) :
+  ∃ s back, to_slice_mut α n v = ret (s, back) ∧
+  v.val = s.val ∧
+  back = Array.from_slice α n v
+  := by simp [to_slice_mut, to_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