cleanup: split TreeKEM.remove into several sub-functions

Inria-Prosecco · Nov 27, 2024 · ab03f26 · ab03f26
1 parent aa1c4a3
commit ab03f26
Showing 1 changed file with 30 additions and 15 deletions.
diff --git a/fstar/treekem/code/MLS.TreeKEM.API.Tree.fst b/fstar/treekem/code/MLS.TreeKEM.API.Tree.fst
@@ -98,40 +98,55 @@ let update #bytes #cb #leaf_ind st lp i =
 
 (*** Remove ***)
 
+val truncate_one:
+  #bytes:Type0 -> {|crypto_bytes bytes|} -> #leaf_ind:nat ->
+  st:treekem_tree_state bytes leaf_ind{1 <= st.levels && is_tree_empty (TNode?.right st.tree)} ->
+  res:treekem_tree_state bytes leaf_ind{res.levels == st.levels-1}
+let truncate_one #bytes #cb #leaf_ind st =
+  if leaf_ind >= pow2 (st.levels-1) then (
+    MLS.TreeCommon.Lemmas.is_tree_empty_leaf_at (TNode?.right st.tree) leaf_ind;
+    false_elim ()
+  ) else (
+    treekem_invariant_truncate st.tree;
+    treekem_priv_invariant_truncate st.tree st.priv;
+    {
+      levels = st.levels-1;
+      tree = tree_truncate st.tree;
+      priv = path_truncate st.priv;
+    }
+  )
+
 val fully_truncate:
   #bytes:Type0 -> {|crypto_bytes bytes|} -> #leaf_ind:nat ->
   st:treekem_tree_state bytes leaf_ind ->
   Tot (treekem_tree_state bytes leaf_ind)
   (decreases st.levels)
 let rec fully_truncate #bytes #cb #leaf_ind st =
   if 1 <= st.levels && is_tree_empty (TNode?.right st.tree) then (
-    if leaf_ind >= pow2 (st.levels-1) then (
-      MLS.TreeCommon.Lemmas.is_tree_empty_leaf_at (TNode?.right st.tree) leaf_ind;
-      false_elim ()
-    ) else (
-      treekem_invariant_truncate st.tree;
-      treekem_priv_invariant_truncate st.tree st.priv;
-      fully_truncate {
-        levels = st.levels-1;
-        tree = tree_truncate st.tree;
-        priv = path_truncate st.priv;
-      }
-    )
+    fully_truncate (truncate_one st)
   ) else (
     st
   )
 
-val remove:
+val remove_aux:
   #bytes:Type0 -> {|crypto_bytes bytes|} -> #leaf_ind:nat ->
   st:treekem_tree_state bytes leaf_ind -> i:treekem_index st{i <> leaf_ind} ->
   treekem_tree_state bytes leaf_ind
-let remove #bytes #cb #leaf_ind st i =
+let remove_aux #bytes #cb #leaf_ind st i =
   treekem_invariant_remove st.tree i;
   treekem_priv_invariant_remove st.tree st.priv i;
-  fully_truncate { st with
+  { st with
     tree = tree_remove st.tree i;
+    priv = path_blank st.priv i;
   }
 
+val remove:
+  #bytes:Type0 -> {|crypto_bytes bytes|} -> #leaf_ind:nat ->
+  st:treekem_tree_state bytes leaf_ind -> i:treekem_index st{i <> leaf_ind} ->
+  treekem_tree_state bytes leaf_ind
+let remove #bytes #cb #leaf_ind st i =
+  fully_truncate (remove_aux st i)
+
 (*** Process Commit ***)
 
 val commit: