zipper_algebra: pass a more reasonable child mask

pathmap-derive/src/lib.rs

@@ -8,6 +8,7 @@ use std::collections::BTreeSet;
 enum PolyZipperTrait {
     Zipper,
     ZipperValues,
+    ZipperValuesAt,
     ZipperReadOnlyValues,
     ZipperReadOnlyConditionalValues,
     ZipperReadOnlyConditionalIteration,
@@ -26,6 +27,7 @@ impl PolyZipperTrait {
         match ident.to_string().as_str() {
             "Zipper" => Some(Self::Zipper),
             "ZipperValues" => Some(Self::ZipperValues),
+            "ZipperValuesAt" => Some(Self::ZipperValuesAt),
             "ZipperReadOnlyValues" => Some(Self::ZipperReadOnlyValues),
             "ZipperReadOnlyConditionalValues" => Some(Self::ZipperReadOnlyConditionalValues),
             "ZipperReadOnlyConditionalIteration" => Some(Self::ZipperReadOnlyConditionalIteration),
@@ -47,6 +49,7 @@ fn all_poly_zipper_traits() -> BTreeSet<PolyZipperTrait> {
     BTreeSet::from([
         Zipper,
         ZipperValues,
+        ZipperValuesAt,
         ZipperReadOnlyValues,
         ZipperReadOnlyConditionalValues,
         ZipperReadOnlyConditionalIteration,
@@ -118,6 +121,11 @@ fn add_trait_dependencies(traits: &mut BTreeSet<PolyZipperTrait>) {
             }
         }
         if traits.contains(&ZipperInfallibleSubtries) {
+            if traits.insert(ZipperValuesAt) {
+                changed = true;
+            }
+        }
+        if traits.contains(&ZipperValuesAt) {
             if traits.insert(ZipperValues) {
                 changed = true;
             }
@@ -312,7 +320,28 @@ fn derive_poly_zipper_with_traits(
                         #(#variant_arms => inner.val(),)*
                     }
                 }
+            }
+        })
+    } else {
+        None
+    };
 
+    // Generate ZipperValuesAt trait implementation
+    let zipper_values_at_impl = if traits.contains(&PolyZipperTrait::ZipperValuesAt) {
+        let variant_arms = &variant_arms;
+        let zipper_values_where = if include_where_clause {
+            quote! {
+                where
+                    #(#inner_types: pathmap::zipper::ZipperValuesAt<V>,)*
+                    #where_clause
+            }
+        } else {
+            quote! {}
+        };
+        Some(quote! {
+            impl #impl_generics pathmap::zipper::ZipperValuesAt<V> for #enum_name #ty_generics
+            #zipper_values_where
+            {
                 fn val_at<K: AsRef<[u8]>>(&self, path: K) -> Option<&V> {
                     match self {
                         #(#variant_arms => inner.val_at(path),)*
@@ -825,6 +854,7 @@ fn derive_poly_zipper_with_traits(
         #(#from_impls)*
         #zipper_impl
         #zipper_values_impl
+        #zipper_values_at_impl
         #zipper_read_only_values_impl
         #zipper_read_only_conditional_values_impl
         // #zipper_forking_impl

src/arena_compact.rs

@@ -82,6 +82,7 @@ use std::marker::PhantomData;
 use fast_slice_utils::starts_with;
 
 use crate::alloc::{GlobalAlloc, global_alloc};
+use crate::zipper::ZipperValuesAt;
 use crate::{
     PathMap,
     morphisms::Catamorphism,
@@ -2785,6 +2786,11 @@ where Storage: AsRef<[u8]>
     fn val(&self) -> Option<&()> {
         self.get_value().map(|_x| &())
     }
+}
+
+impl<'tree, Storage> ZipperValuesAt<()> for ACTZipper<'tree, Storage, ()>
+where Storage: AsRef<[u8]>
+{
     fn val_at<K: AsRef<[u8]>>(&self, path: K) -> Option<&()> {
         self.get_value_at(path.as_ref()).map(|_x| &())
     }
@@ -2797,6 +2803,11 @@ where Storage: AsRef<[u8]>
         //GOAT, see soundness discussion in ZipperReadOnlyValues impl below
         self.get_val()
     }
+}
+
+impl<'tree, Storage> ZipperValuesAt<u64> for ACTZipper<'tree, Storage, u64>
+where Storage: AsRef<[u8]>
+{
     fn val_at<K: AsRef<[u8]>>(&self, path: K) -> Option<&u64> {
         //GOAT, see soundness discussion in ZipperReadOnlyValues impl below
         self.get_val_at(path)

src/dependent_zipper.rs

@@ -237,6 +237,15 @@ impl<'trie, PrimaryZ, SecondaryZ, V, C, F : Clone + for <'a> FnOnce(C, &'a [u8],
             self.primary.val()
         }
     }
+}
+
+impl<'trie, PrimaryZ, SecondaryZ, V, C, F : Clone + for <'a> FnOnce(C, &'a [u8], usize) -> (C, Option<SecondaryZ>)> ZipperValuesAt<V>
+    for DependentProductZipperG<'trie, PrimaryZ, SecondaryZ, V, C, F>
+    where
+        V: Clone + Send + Sync,
+        PrimaryZ: ZipperMoving + ZipperValuesAt<V>,
+        SecondaryZ: ZipperMoving + ZipperValuesAt<V>,
+{
     fn val_at<K: AsRef<[u8]>>(&self, path: K) -> Option<&V> {
         if let Some(idx) = self.factor_idx(true) {
             self.secondary[idx].val_at(path)

src/empty_zipper.rs

@@ -90,6 +90,9 @@ impl ZipperIteration for EmptyZipper {
 
 impl<V> ZipperValues<V> for EmptyZipper {
     fn val(&self) -> Option<&V> { None }
+}
+
+impl<V> ZipperValuesAt<V> for EmptyZipper {
     fn val_at<K: AsRef<[u8]>>(&self, _path: K) -> Option<&V> { None }
 }
 

src/experimental/zipper_algebra.rs

@@ -515,11 +515,11 @@ where
         Out: ZipperWriting<V, A>,
     {
         if *lhs_grafts != ByteMask::EMPTY {
-            out.graft_masked_branches(lhs, std::mem::take(lhs_grafts), false);
+            out.graft_masked_branches(lhs, std::mem::take(lhs_grafts) & lhs.child_mask(), false);
         }
 
         if *rhs_grafts != ByteMask::EMPTY {
-            out.graft_masked_branches(rhs, std::mem::take(rhs_grafts), false);
+            out.graft_masked_branches(rhs, std::mem::take(rhs_grafts) & rhs.child_mask(), false);
         }
     }
 
@@ -729,15 +729,15 @@ where
         Out: ZipperWriting<V, A>,
     {
         if *lhs_grafts != ByteMask::EMPTY {
-            out.graft_masked_branches(lhs, std::mem::take(lhs_grafts), false);
+            out.graft_masked_branches(lhs, std::mem::take(lhs_grafts) & lhs.child_mask(), false);
         }
 
         if *mid_grafts != ByteMask::EMPTY {
-            out.graft_masked_branches(mid, std::mem::take(mid_grafts), false);
+            out.graft_masked_branches(mid, std::mem::take(mid_grafts) & mid.child_mask(), false);
         }
 
         if *rhs_grafts != ByteMask::EMPTY {
-            out.graft_masked_branches(rhs, std::mem::take(rhs_grafts), false);
+            out.graft_masked_branches(rhs, std::mem::take(rhs_grafts) & rhs.child_mask(), false);
         }
     }
 
@@ -1011,19 +1011,19 @@ fn zipper_merge4<P, V, Z0, Z1, Z2, Z3, Out, A>(
         Out: ZipperWriting<V, A>,
     {
         if *z0_grafts != ByteMask::EMPTY {
-            out.graft_masked_branches(z0, std::mem::take(z0_grafts), false);
+            out.graft_masked_branches(z0, std::mem::take(z0_grafts) & z0.child_mask(), false);
         }
 
         if *z1_grafts != ByteMask::EMPTY {
-            out.graft_masked_branches(z1, std::mem::take(z1_grafts), false);
+            out.graft_masked_branches(z1, std::mem::take(z1_grafts) & z1.child_mask(), false);
         }
 
         if *z2_grafts != ByteMask::EMPTY {
-            out.graft_masked_branches(z2, std::mem::take(z2_grafts), false);
+            out.graft_masked_branches(z2, std::mem::take(z2_grafts) & z2.child_mask(), false);
         }
 
         if *z3_grafts != ByteMask::EMPTY {
-            out.graft_masked_branches(z3, std::mem::take(z3_grafts), false);
+            out.graft_masked_branches(z3, std::mem::take(z3_grafts) & z3.child_mask(), false);
         }
     }
 
@@ -1665,7 +1665,12 @@ where
     {
         for_each_bit(active, |i| {
             if grafts[i] != ByteMask::EMPTY {
-                out.graft_masked_branches(&zs[i], std::mem::take(&mut grafts[i]), false);
+                let z = &zs[i];
+                out.graft_masked_branches(
+                    z,
+                    std::mem::take(&mut grafts[i]) & z.child_mask(),
+                    false,
+                );
             }
         });
     }
@@ -2856,7 +2861,7 @@ mod zipper_algebra_poly {
     use pathmap_derive::PolyZipperExplicit;
 
     #[derive(PolyZipperExplicit)]
-    #[poly_zipper_explicit(traits(ZipperMoving, ZipperValues, ZipperConcrete))]
+    #[poly_zipper_explicit(traits(ZipperMoving, ZipperValues, ZipperValuesAt, ZipperConcrete))]
     pub(super) enum SomeMutRefZ<'a, 'trie, 'path, V: Clone + Send + Sync + Unpin, A: Allocator> {
         RZ(&'a mut ReadZipperUntracked<'trie, 'path, V, A>),
         RZT(&'a mut ReadZipperTracked<'trie, 'path, V, A>),

src/overlay_zipper.rs

@@ -18,7 +18,7 @@
 
 use fast_slice_utils::find_prefix_overlap;
 use crate::utils::{BitMask, ByteMask};
-use crate::zipper::{Zipper, ZipperMoving, ZipperIteration, ZipperValues};
+use crate::zipper::{Zipper, ZipperMoving, ZipperIteration, ZipperValues, ZipperValuesAt};
 
 /// Zipper that traverses a virtual trie formed by fusing the tries of two other zippers
 pub struct OverlayZipper<AV, BV, OutV, AZipper, BZipper, Mapping>
@@ -104,6 +104,15 @@ impl<AV, BV, OutV, AZipper, BZipper, Mapping> ZipperValues<OutV>
     fn val(&self) -> Option<&OutV> {
         (self.mapping)(self.a.val(), self.b.val())
     }
+}
+
+impl<AV, BV, OutV, AZipper, BZipper, Mapping> ZipperValuesAt<OutV>
+    for OverlayZipper<AV, BV, OutV, AZipper, BZipper, Mapping>
+    where
+        AZipper: ZipperValuesAt<AV>,
+        BZipper: ZipperValuesAt<BV>,
+        Mapping: for<'a> Fn(Option<&'a AV>, Option<&'a BV>) -> Option<&'a OutV>,
+{
     fn val_at<K: AsRef<[u8]>>(&self, path: K) -> Option<&OutV> {
         (self.mapping)(self.a.val_at(&path), self.b.val_at(&path))
     }

src/poly_zipper.rs

@@ -142,7 +142,7 @@ mod tests {
     // ======================================================================================
     // Cocktail of recursive zipper madness
     #[derive(PolyZipperExplicit)]
-    #[poly_zipper_explicit(traits(Zipper, ZipperValues, ZipperMoving, ZipperIteration))]
+    #[poly_zipper_explicit(traits(Zipper, ZipperValues, ZipperValuesAt, ZipperMoving, ZipperIteration))]
     pub enum ExprFactor<'trie, V: Clone + Send + Sync + Unpin + 'static = ()> {
         Specific(ReadZipperOwned<V>),
         Generic(PrefixZipper<'trie,

src/prefix_zipper.rs

@@ -242,6 +242,12 @@ impl<'prefix, Z, V> ZipperValues<V> for PrefixZipper<'prefix, Z>
         }
         self.source.val()
     }
+}
+
+impl<'prefix, Z, V> ZipperValuesAt<V> for PrefixZipper<'prefix, Z>
+    where
+        Z: ZipperValuesAt<V>
+{
     fn val_at<K: AsRef<[u8]>>(&self, path: K) -> Option<&V> {
         let path = self.adjust_lookup_path(path.as_ref())?;
         self.source.val_at(path)
@@ -687,6 +693,7 @@ mod tests {
     use crate::zipper::ZipperAbsolutePath;
     use crate::zipper::ZipperReadOnlyValues;
     use crate::zipper::ZipperValues;
+    use crate::zipper::ZipperValuesAt;
     const PATHS1: &[(&[u8], u64)] = &[
         (b"0000", 0),
         (b"00000", 1),

src/product_zipper.rs

@@ -308,6 +308,9 @@ impl<'trie, V: Clone + Send + Sync + Unpin + 'trie, A: Allocator + 'trie> Zipper
     fn val(&self) -> Option<&V> {
         unsafe{ self.z.get_val() }
     }
+}
+
+impl<'trie, V: Clone + Send + Sync + Unpin + 'trie, A: Allocator + 'trie> ZipperValuesAt<V> for ProductZipper<'_, 'trie, V, A> {
     fn val_at<K: AsRef<[u8]>>(&self, path: K) -> Option<&V> {
         unsafe{ self.z.get_val_at(path) }
     }
@@ -562,6 +565,15 @@ impl<'trie, PrimaryZ, SecondaryZ, V> ZipperValues<V>
             self.primary.val()
         }
     }
+}
+
+impl<'trie, PrimaryZ, SecondaryZ, V> ZipperValuesAt<V>
+    for ProductZipperG<'trie, PrimaryZ, SecondaryZ, V>
+    where
+        V: Clone + Send + Sync,
+        PrimaryZ: ZipperMoving + ZipperValuesAt<V>,
+        SecondaryZ: ZipperMoving + ZipperValuesAt<V>,
+{
     fn val_at<K: AsRef<[u8]>>(&self, path: K) -> Option<&V> {
         if let Some(idx) = self.factor_idx(true) {
             self.secondary[idx].val_at(path)
@@ -887,6 +899,7 @@ impl <Z : ZipperAbsolutePath> ZipperAbsolutePath for OneFactor<Z> { zipper_impl_
 impl <Z : ZipperMoving> ZipperMoving for OneFactor<Z> { zipper_impl_lens!(ZipperMoving self => self.z); }
 impl <Z : ZipperIteration> ZipperIteration for OneFactor<Z> { zipper_impl_lens!(ZipperIteration self => self.z); }
 impl <V, Z : ZipperValues<V>> ZipperValues<V> for OneFactor<Z> { zipper_impl_lens!(ZipperValues self => self.z); }
+impl <V, Z : ZipperValuesAt<V>> ZipperValuesAt<V> for OneFactor<Z> { zipper_impl_lens!(ZipperValuesAt self => self.z); }
 impl <V, Z : ZipperForking<V>> ZipperForking<V> for OneFactor<Z> { type ReadZipperT<'a> = Z::ReadZipperT<'a> where Z: 'a; zipper_impl_lens!(ZipperForking self => self.z); }
 impl <V: Clone + Send + Sync, A: Allocator, Z : ZipperSubtries<V, A>> ZipperSubtries<V, A> for OneFactor<Z> { zipper_impl_lens!(ZipperSubtries self => self.z); }
 impl <V: Clone + Send + Sync, A: Allocator, Z : ZipperInfallibleSubtries<V, A>> ZipperInfallibleSubtries<V, A> for OneFactor<Z> { zipper_impl_lens!(ZipperInfallibleSubtries self => self.z); }

src/trie_ref.rs

@@ -266,6 +266,9 @@ impl<V: Clone + Send + Sync + Unpin, A: Allocator> ZipperValues<V> for TrieRefBo
     fn val(&self) -> Option<&V> {
         self.get_val()
     }
+}
+
+impl<V: Clone + Send + Sync + Unpin, A: Allocator> ZipperValuesAt<V> for TrieRefBorrowed<'_, V, A> {
     fn val_at<K: AsRef<[u8]>>(&self, path: K) -> Option<&V> {
         self.get_val_at(path)
     }
@@ -646,6 +649,9 @@ impl<V: Clone + Send + Sync + Unpin, A: Allocator> ZipperValues<V> for TrieRefOw
             None
         }
     }
+}
+
+impl<V: Clone + Send + Sync + Unpin, A: Allocator> ZipperValuesAt<V> for TrieRefOwned<V, A> {
     fn val_at<K: AsRef<[u8]>>(&self, path: K) -> Option<&V> {
         if self.is_valid() {
             TrieRefBorrowed::new_with_key_and_path_in(
@@ -844,6 +850,9 @@ impl<V: Clone + Send + Sync + Unpin, A: Allocator> ZipperValues<V> for TrieRef<'
             TrieRef::Owned(trie_ref) => trie_ref.val(),
         }
     }
+}
+
+impl<V: Clone + Send + Sync + Unpin, A: Allocator> ZipperValuesAt<V> for TrieRef<'_, V, A> {
     fn val_at<K: AsRef<[u8]>>(&self, path: K) -> Option<&V> {
         match self {
             TrieRef::Borrowed(trie_ref) => trie_ref.val_at(path),
@@ -1138,7 +1147,7 @@ mod tests {
 
     #[test]
     fn trie_ref_val_at_test() {
-        fn assert_val_at<T: ZipperValues<i32>>(trie_ref: T) {
+        fn assert_val_at<T: ZipperValuesAt<i32>>(trie_ref: T) {
             assert_eq!(trie_ref.val(), None);
             assert_eq!(trie_ref.val_at(b"root:a:new_a"), Some(&10));
             assert_eq!(trie_ref.val_at(b"root:a:nested:deep"), Some(&11));

src/write_zipper.rs

@@ -390,6 +390,9 @@ impl<'a, V: Clone + Send + Sync + Unpin, A: Allocator + 'a> Zipper for WriteZipp
 
 impl<'a, V: Clone + Send + Sync + Unpin, A: Allocator + 'a> ZipperValues<V> for WriteZipperTracked<'a, '_, V, A>{
     fn val(&self) -> Option<&V> { self.z.val() }
+}
+
+impl<'a, V: Clone + Send + Sync + Unpin, A: Allocator + 'a> ZipperValuesAt<V> for WriteZipperTracked<'a, '_, V, A>{
     fn val_at<K: AsRef<[u8]>>(&self, path: K) -> Option<&V> { self.z.val_at(path) }
 }
 
@@ -550,6 +553,9 @@ impl<'a, V: Clone + Send + Sync + Unpin, A: Allocator + 'a> Zipper for WriteZipp
 
 impl<'a, V: Clone + Send + Sync + Unpin, A: Allocator + 'a> ZipperValues<V> for WriteZipperUntracked<'a, '_, V, A> {
     fn val(&self) -> Option<&V> { self.z.val() }
+}
+
+impl<'a, V: Clone + Send + Sync + Unpin, A: Allocator + 'a> ZipperValuesAt<V> for WriteZipperUntracked<'a, '_, V, A> {
     fn val_at<K: AsRef<[u8]>>(&self, path: K) -> Option<&V> { self.z.val_at(path) }
 }
 
@@ -717,6 +723,7 @@ impl<V: 'static + Clone + Send + Sync + Unpin, A: Allocator> Clone for WriteZipp
 
 impl<V: Clone + Send + Sync + Unpin, A: Allocator> Zipper for WriteZipperOwned<V, A> { zipper_impl_lens!(Zipper self => self.z); }
 impl<V: Clone + Send + Sync + Unpin, A: Allocator> ZipperValues<V> for WriteZipperOwned<V, A> { zipper_impl_lens!(ZipperValues self => self.z); }
+impl<V: Clone + Send + Sync + Unpin, A: Allocator> ZipperValuesAt<V> for WriteZipperOwned<V, A> { zipper_impl_lens!(ZipperValuesAt self => self.z); }
 impl<V: Clone + Send + Sync + Unpin, A: Allocator> ZipperInfallibleSubtries<V, A> for WriteZipperOwned<V, A> { zipper_impl_lens!(ZipperInfallibleSubtries self => self.z); }
 impl<V: Clone + Send + Sync + Unpin, A: Allocator> ZipperMoving for WriteZipperOwned<V, A> { zipper_impl_lens!(ZipperMoving self => self.z); }
 impl<V: Clone + Send + Sync + Unpin, A: Allocator> ZipperPathBuffer for WriteZipperOwned<V, A> { zipper_impl_lens!(ZipperPathBuffer self => self.z); }