chore: use RawPtr::Env to signal raw pointers for environments

It's safer to use a new variant `RawPtr::Env` to signal raw pointers
for environments than relying solely on tag checks. This would
probably have avoided the bug introduced in #1183, whose fix was
implemented in #1200

benches/common/fib.rs

@@ -58,8 +58,8 @@ fn lurk_fib<F: LurkField>(store: &Store<F>, n: usize) -> Ptr {
     //               saved_env: (((.lurk.user.next . <FUNCTION (.lurk.user.a .lurk.user.b) (.lurk.user.next .lurk.user.b (+ .lurk.user.a .lurk.user.b))>))),
     //               body: (.lurk.user.fib), continuation: Outermost }
 
-    let [_, _, rest_bindings] = store.pop_binding(*target_env).unwrap();
-    let [_, val, _] = store.pop_binding(rest_bindings).unwrap();
+    let [_, _, rest_bindings] = store.pop_binding(target_env).unwrap();
+    let [_, val, _] = store.pop_binding(&rest_bindings).unwrap();
     val
 }
 

foil/src/coil.rs

@@ -81,7 +81,7 @@ pub struct Context {
 /// Look up `var` in `env`, where `env` is a list of bindings from `Symbol` to `U64` representing bindings of variables to vertices
 /// by id.
 fn lookup_vertex_id<F: LurkField>(store: &Store<F>, env: &Ptr, var: &Ptr) -> Result<Option<Id>> {
-    let Some([bound_var, id, rest_env]) = store.pop_binding(*env) else {
+    let Some([bound_var, id, rest_env]) = store.pop_binding(env) else {
         return Ok(None);
     };
 
@@ -125,7 +125,7 @@ impl Context {
 
     fn pop_binding<F: LurkField>(&mut self, store: &Store<F>) -> Result<Id> {
         let [_var, id, rest_env] = store
-            .pop_binding(self.env)
+            .pop_binding(&self.env)
             .ok_or(anyhow!("failed to pop binding"))?;
         self.env = rest_env;
 

src/cli/zstore.rs

@@ -51,38 +51,20 @@ impl<F: LurkField> ZDag<F> {
                         z_ptr
                     }
                     RawPtr::Hash4(idx) => {
-                        if matches!(tag, Tag::Expr(Env)) {
-                            let [sym, val, env] = store.expect_env_components(*idx);
-                            let sym = self.populate_with(&sym, store, cache);
-                            let val = self.populate_with(&val, store, cache);
-                            let env = self.populate_with(&env, store, cache);
-                            let z_ptr = ZPtr::from_parts(
-                                *tag,
-                                store.poseidon_cache.hash4(&[
-                                    *sym.value(),
-                                    val.tag_field(),
-                                    *val.value(),
-                                    *env.value(),
-                                ]),
-                            );
-                            self.0.insert(z_ptr, ZPtrType::Env(sym, val, env));
-                            z_ptr
-                        } else {
-                            let [a, b] = expect_ptrs!(store, 2, *idx);
-                            let a = self.populate_with(&a, store, cache);
-                            let b = self.populate_with(&b, store, cache);
-                            let z_ptr = ZPtr::from_parts(
-                                *tag,
-                                store.poseidon_cache.hash4(&[
-                                    a.tag_field(),
-                                    *a.value(),
-                                    b.tag_field(),
-                                    *b.value(),
-                                ]),
-                            );
-                            self.0.insert(z_ptr, ZPtrType::Tuple2(a, b));
-                            z_ptr
-                        }
+                        let [a, b] = expect_ptrs!(store, 2, *idx);
+                        let a = self.populate_with(&a, store, cache);
+                        let b = self.populate_with(&b, store, cache);
+                        let z_ptr = ZPtr::from_parts(
+                            *tag,
+                            store.poseidon_cache.hash4(&[
+                                a.tag_field(),
+                                *a.value(),
+                                b.tag_field(),
+                                *b.value(),
+                            ]),
+                        );
+                        self.0.insert(z_ptr, ZPtrType::Tuple2(a, b));
+                        z_ptr
                     }
                     RawPtr::Hash6(idx) => {
                         let [a, b, c] = expect_ptrs!(store, 3, *idx);
@@ -125,6 +107,23 @@ impl<F: LurkField> ZDag<F> {
                         self.0.insert(z_ptr, ZPtrType::Tuple4(a, b, c, d));
                         z_ptr
                     }
+                    RawPtr::Env(idx) => {
+                        let [sym, val, env] = store.expect_env_ptrs(*idx);
+                        let sym = self.populate_with(&sym, store, cache);
+                        let val = self.populate_with(&val, store, cache);
+                        let env = self.populate_with(&env, store, cache);
+                        let z_ptr = ZPtr::from_parts(
+                            *tag,
+                            store.poseidon_cache.hash4(&[
+                                *sym.value(),
+                                val.tag_field(),
+                                *val.value(),
+                                *env.value(),
+                            ]),
+                        );
+                        self.0.insert(z_ptr, ZPtrType::Env(sym, val, env));
+                        z_ptr
+                    }
                 };
                 cache.insert(*ptr, z_ptr);
                 z_ptr
@@ -134,7 +133,7 @@ impl<F: LurkField> ZDag<F> {
         let mut stack = vec![*ptr];
         macro_rules! feed_loop {
             ($x:expr) => {
-                if $x.raw().is_hash() {
+                if $x.raw().is_not_atom() {
                     if !cache.contains_key(&$x) {
                         if dag.insert($x) {
                             stack.push($x);
@@ -147,14 +146,8 @@ impl<F: LurkField> ZDag<F> {
             match ptr.raw() {
                 RawPtr::Atom(..) => (),
                 RawPtr::Hash4(idx) => {
-                    if matches!(ptr.tag(), Tag::Expr(Env)) {
-                        for ptr in store.expect_env_components(*idx) {
-                            feed_loop!(ptr)
-                        }
-                    } else {
-                        for ptr in expect_ptrs!(store, 2, *idx) {
-                            feed_loop!(ptr)
-                        }
+                    for ptr in expect_ptrs!(store, 2, *idx) {
+                        feed_loop!(ptr)
                     }
                 }
                 RawPtr::Hash6(idx) => {
@@ -167,6 +160,11 @@ impl<F: LurkField> ZDag<F> {
                         feed_loop!(ptr)
                     }
                 }
+                RawPtr::Env(idx) => {
+                    for ptr in store.expect_env_ptrs(*idx) {
+                        feed_loop!(ptr)
+                    }
+                }
             }
         }
         dag.iter()
@@ -223,15 +221,14 @@ impl<F: LurkField> ZDag<F> {
                         intern_ptrs_hydrated!(store, *z_ptr.tag(), *z_ptr, ptr1, ptr2, ptr3, ptr4)
                     }
                     Some(ZPtrType::Env(sym, val, env)) => {
-                        let sym = self.populate_store(sym, store, cache)?;
-                        let val = self.populate_store(val, store, cache)?;
-                        let env = self.populate_store(env, store, cache)?;
-                        let raw = store.intern_raw_ptrs([
-                            *sym.raw(),
-                            store.tag(*val.tag()),
-                            *val.raw(),
-                            *env.raw(),
-                        ]);
+                        let (_, sym_raw) = self.populate_store(sym, store, cache)?.into_parts();
+                        let (val_tag, val_raw) =
+                            self.populate_store(val, store, cache)?.into_parts();
+                        let (_, env_raw) = self.populate_store(env, store, cache)?.into_parts();
+                        let raw = store.intern_raw_ptrs_hydrated(
+                            [sym_raw, store.tag(val_tag), val_raw, env_raw],
+                            *z_ptr.value(),
+                        );
                         Ptr::new(Tag::Expr(Env), raw)
                     }
                 };

src/coprocessor/gadgets.rs

@@ -215,7 +215,7 @@ pub(crate) fn deconstruct_env<F: LurkField, CS: ConstraintSystem<F>>(
     let env_ptr = s.to_ptr(&env_zptr);
 
     let (a, b, c, d) = {
-        if let Some([v, val, new_env]) = s.pop_binding(env_ptr) {
+        if let Some([v, val, new_env]) = s.pop_binding(&env_ptr) {
             let v_zptr = s.hash_ptr(&v);
             let val_zptr = s.hash_ptr(&val);
             let new_env_zptr = s.hash_ptr(&new_env);

src/coroutine/memoset/env.rs

@@ -32,7 +32,7 @@ impl<F: LurkField> Query<F> for EnvQuery<F> {
         let s = scope.store.as_ref();
         match self {
             Self::Lookup(var, env) => {
-                if let Some([v, val, new_env]) = s.pop_binding(*env) {
+                if let Some([v, val, new_env]) = s.pop_binding(env) {
                     if s.ptr_eq(var, &v) {
                         let t = s.intern_t();
                         s.cons(val, t)

src/lem/interpreter.rs

@@ -422,7 +422,7 @@ impl Block {
                 Op::PopBinding(preimg, img) => {
                     let img_ptr = bindings.get_ptr(img)?;
                     let preimg_ptrs = store
-                        .pop_binding(img_ptr)
+                        .pop_binding(&img_ptr)
                         .context("cannot extract {img}'s binding")?;
                     for (var, ptr) in preimg.iter().zip(preimg_ptrs.iter()) {
                         bindings.insert_ptr(var.clone(), *ptr);

src/lem/pointers.rs

@@ -10,17 +10,25 @@ use super::Tag;
 /// `RawPtr` is the basic pointer type of the LEM store. An `Atom` points to a field
 /// element, and a `HashN` points to `N` children, which are also raw pointers. Thus,
 /// they are a building block for graphs that represent Lurk data.
+///
+/// The `RawPtr::Env` variant has 4 children and is meant to encode environments in
+/// a more compact way than cons lists. The 4 children, in order, are:
+/// 1. The `raw` component of the symbol from the head binding
+/// 2. The `tag` component of the value from the head binding
+/// 3. The `raw` component of the value from the head binding
+/// 4. The `raw` component of the tail of the environment
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize, Hash)]
 pub enum RawPtr {
     Atom(usize),
     Hash4(usize),
     Hash6(usize),
     Hash8(usize),
+    Env(usize),
 }
 
 impl RawPtr {
     #[inline]
-    pub fn is_hash(&self) -> bool {
+    pub fn is_not_atom(&self) -> bool {
         matches!(
             self,
             RawPtr::Hash4(..) | RawPtr::Hash6(..) | RawPtr::Hash8(..)
@@ -58,6 +66,14 @@ impl RawPtr {
             _ => None,
         }
     }
+
+    #[inline]
+    pub fn get_env(&self) -> Option<usize> {
+        match self {
+            RawPtr::Env(x) => Some(*x),
+            _ => None,
+        }
+    }
 }
 
 /// `Ptr` is a tagged pointer. The tag is there to say what kind of data it encodes.
@@ -91,6 +107,12 @@ impl Ptr {
         (tag, raw)
     }
 
+    #[inline]
+    pub fn into_parts(self) -> (Tag, RawPtr) {
+        let Ptr { tag, raw } = self;
+        (tag, raw)
+    }
+
     #[inline]
     pub fn has_tag(&self, tag: &Tag) -> bool {
         self.tag() == tag
@@ -156,6 +178,11 @@ impl Ptr {
         self.raw().get_hash8()
     }
 
+    #[inline]
+    pub fn get_env(&self) -> Option<usize> {
+        self.raw().get_env()
+    }
+
     #[inline]
     pub fn atom(tag: Tag, idx: usize) -> Ptr {
         Ptr {

src/lem/store.rs

@@ -281,18 +281,14 @@ impl<F: LurkField> Store<F> {
     /// Similar to `intern_raw_ptrs` but doesn't add the resulting pointer to
     /// `dehydrated`. This function is used when converting a `ZStore` to a
     /// `Store`.
-    pub fn intern_raw_ptrs_hydrated<const N: usize>(
-        &self,
-        ptrs: [RawPtr; N],
-        z: FWrap<F>,
-    ) -> RawPtr {
+    pub fn intern_raw_ptrs_hydrated<const N: usize>(&self, ptrs: [RawPtr; N], z: F) -> RawPtr {
         let (ptr, _) = self.intern_raw_ptrs_internal::<N>(ptrs);
-        self.z_cache.insert(ptr, Box::new(z));
-        self.inverse_z_cache.insert(z, Box::new(ptr));
+        let z_wrap = FWrap(z);
+        self.z_cache.insert(ptr, Box::new(z_wrap));
+        self.inverse_z_cache.insert(z_wrap, Box::new(ptr));
         ptr
     }
 
-    #[inline]
     fn intern_raw_ptrs_internal<const N: usize>(&self, ptrs: [RawPtr; N]) -> (RawPtr, bool) {
         macro_rules! intern {
             ($Hash:ident, $hash:ident, $n:expr) => {{
@@ -326,7 +322,7 @@ impl<F: LurkField> Store<F> {
         z: ZPtr<F>,
     ) -> Ptr {
         let raw_ptrs = self.ptrs_to_raw_ptrs::<N, P>(&ptrs);
-        let payload = self.intern_raw_ptrs_hydrated::<N>(raw_ptrs, FWrap(*z.value()));
+        let payload = self.intern_raw_ptrs_hydrated::<N>(raw_ptrs, *z.value());
         Ptr::new(tag, payload)
     }
 
@@ -391,28 +387,38 @@ impl<F: LurkField> Store<F> {
 
     #[inline]
     pub fn fetch_tag(&self, ptr: &RawPtr) -> Option<Tag> {
-        let idx = ptr.get_atom()?;
-        Tag::pos(idx)
+        ptr.get_atom().and_then(Tag::pos)
     }
 
     pub fn raw_to_ptr(&self, tag: &RawPtr, raw: &RawPtr) -> Option<Ptr> {
         let tag = self.fetch_tag(tag)?;
         Some(Ptr::new(tag, *raw))
     }
 
-    #[inline]
+    /// Creates a `RawPtr` for an environment
+    pub fn intern_env_raw_components(&self, ptrs: [RawPtr; 4]) -> RawPtr {
+        let (idx, inserted) = self.hash4.insert_probe(Box::new(ptrs));
+        let ptr = RawPtr::Env(idx);
+        if inserted {
+            // this is for `hydrate_z_cache`
+            self.dehydrated.load().push(Box::new(ptr));
+        }
+        ptr
+    }
+
     pub fn push_binding(&self, sym: Ptr, val: Ptr, env: Ptr) -> Ptr {
-        assert_eq!(*sym.tag(), Tag::Expr(Sym));
-        assert_eq!(*env.tag(), Tag::Expr(Env));
-        let raw =
-            self.intern_raw_ptrs::<4>([*sym.raw(), self.tag(*val.tag()), *val.raw(), *env.raw()]);
+        let (sym_tag, sym_raw) = sym.into_parts();
+        let (val_tag, val_raw) = val.into_parts();
+        let (env_tag, env_raw) = env.into_parts();
+        assert_eq!(sym_tag, Tag::Expr(Sym));
+        assert_eq!(env_tag, Tag::Expr(Env));
+        let raw = self.intern_env_raw_components([sym_raw, self.tag(val_tag), val_raw, env_raw]);
         Ptr::new(Tag::Expr(Env), raw)
     }
 
-    #[inline]
-    pub fn pop_binding(&self, env: Ptr) -> Option<[Ptr; 3]> {
+    pub fn pop_binding(&self, env: &Ptr) -> Option<[Ptr; 3]> {
         assert_eq!(*env.tag(), Tag::Expr(Env));
-        let idx = env.get_index2()?;
+        let idx = env.get_env()?;
         let [sym_pay, val_tag, val_pay, env_pay] = self.fetch_raw_ptrs::<4>(idx)?;
         let val_tag = self.fetch_tag(val_tag)?;
         let sym = Ptr::new(Tag::Expr(Sym), *sym_pay);
@@ -437,7 +443,6 @@ impl<F: LurkField> Store<F> {
         Ptr::new(Tag::Expr(Prov), raw)
     }
 
-    #[inline]
     pub fn deconstruct_provenance(&self, prov: Ptr) -> Option<[Ptr; 3]> {
         assert_eq!(*prov.tag(), Tag::Expr(Prov));
         let idx = prov.get_index2()?;
@@ -881,7 +886,9 @@ impl<F: LurkField> Store<F> {
         }
     }
 
-    pub fn expect_env_components(&self, idx: usize) -> [Ptr; 3] {
+    /// Returns the pointers that form an environment pointer by the index of
+    /// its respective `RawPtr::Env` raw component
+    pub fn expect_env_ptrs(&self, idx: usize) -> [Ptr; 3] {
         let [sym_pay, val_tag, val_pay, env_pay] = self.expect_raw_ptrs(idx);
         let sym = Ptr::new(Tag::Expr(Sym), *sym_pay);
         let val = Ptr::new(
@@ -1021,7 +1028,7 @@ impl<F: LurkField> Store<F> {
         }
         match ptr {
             RawPtr::Atom(idx) => FWrap(*self.expect_f(*idx)),
-            RawPtr::Hash4(idx) => hash_raw!(hash4, 4, *idx),
+            RawPtr::Hash4(idx) | RawPtr::Env(idx) => hash_raw!(hash4, 4, *idx),
             RawPtr::Hash6(idx) => hash_raw!(hash6, 6, *idx),
             RawPtr::Hash8(idx) => hash_raw!(hash8, 8, *idx),
         }
@@ -1074,7 +1081,7 @@ impl<F: LurkField> Store<F> {
         let mut stack = vec![ptr];
         macro_rules! feed_loop {
             ($x:expr) => {
-                if $x.is_hash() {
+                if $x.is_not_atom() {
                     if self.z_cache.get($x).is_none() {
                         if ptrs.insert($x) {
                             stack.push($x);
@@ -1086,7 +1093,7 @@ impl<F: LurkField> Store<F> {
         while let Some(ptr) = stack.pop() {
             match ptr {
                 RawPtr::Atom(..) => (),
-                RawPtr::Hash4(idx) => {
+                RawPtr::Hash4(idx) | RawPtr::Env(idx) => {
                     for ptr in self.expect_raw_ptrs::<4>(*idx) {
                         feed_loop!(ptr)
                     }