From 91465dc78ec7b4a8c9b651657bb8ad5f25c556a7 Mon Sep 17 00:00:00 2001 From: Vincent Ambo Date: Thu, 29 Dec 2022 17:08:14 +0300 Subject: refactor(tvix/eval): persistent, memory-sharing OrdMap for NixAttrs This uses the `im::OrdMap` for `NixAttrs` to enable sharing of memory between different iterations of a map. This slightly speeds up eval, but not significantly. Future work might include benchmarking whether using a `HashMap` and only ordering in cases where order is actually required would help. This switches to a fork of `im` that fixes some bugs with its OrdMap implementation. Change-Id: I2f6a5ff471b6d508c1e8a98b13f889f49c0d9537 Reviewed-on: https://cl.tvl.fyi/c/depot/+/7676 Reviewed-by: sterni Tested-by: BuildkiteCI --- tvix/eval/src/value/attrs.rs | 101 +++++++++++++++++++++---------------------- 1 file changed, 50 insertions(+), 51 deletions(-) (limited to 'tvix/eval/src/value/attrs.rs') diff --git a/tvix/eval/src/value/attrs.rs b/tvix/eval/src/value/attrs.rs index 833cb2fd60..adf56567fb 100644 --- a/tvix/eval/src/value/attrs.rs +++ b/tvix/eval/src/value/attrs.rs @@ -5,10 +5,10 @@ //! //! Due to this, construction and management of attribute sets has //! some peculiarities that are encapsulated within this module. -use std::collections::btree_map; -use std::collections::BTreeMap; use std::iter::FromIterator; +use imbl::{ordmap, OrdMap}; + use crate::errors::ErrorKind; use crate::vm::VM; @@ -23,7 +23,8 @@ mod tests; #[derive(Clone, Debug)] enum AttrsRep { Empty, - Map(BTreeMap), + + Im(OrdMap), /// Warning: this represents a **two**-attribute attrset, with /// attribute names "name" and "value", like `{name="foo"; @@ -43,20 +44,21 @@ impl Default for AttrsRep { impl AttrsRep { /// Retrieve reference to a mutable map inside of an attrs, /// optionally changing the representation if required. - fn map_mut(&mut self) -> &mut BTreeMap { + fn map_mut(&mut self) -> &mut OrdMap { match self { - AttrsRep::Map(m) => m, + AttrsRep::Im(m) => m, AttrsRep::Empty => { - *self = AttrsRep::Map(BTreeMap::new()); + *self = AttrsRep::Im(OrdMap::new()); self.map_mut() } AttrsRep::KV { name, value } => { - *self = AttrsRep::Map(BTreeMap::from([ - (NixString::NAME, name.clone()), - (NixString::VALUE, value.clone()), - ])); + *self = AttrsRep::Im(ordmap! { + NixString::NAME => name.clone(), + NixString::VALUE => value.clone() + }); + self.map_mut() } } @@ -72,7 +74,7 @@ impl AttrsRep { _ => None, }, - AttrsRep::Map(map) => map.get(&key.into()), + AttrsRep::Im(map) => map.get(&key.into()), } } @@ -80,7 +82,7 @@ impl AttrsRep { match self { AttrsRep::Empty => false, AttrsRep::KV { .. } => key == "name" || key == "value", - AttrsRep::Map(map) => map.contains_key(&key.into()), + AttrsRep::Im(map) => map.contains_key(&key.into()), } } } @@ -98,7 +100,7 @@ where where T: IntoIterator, { - NixAttrs(AttrsRep::Map( + NixAttrs(AttrsRep::Im( iter.into_iter() .map(|(k, v)| (k.into(), v.into())) .collect(), @@ -121,7 +123,7 @@ impl TotalDisplay for NixAttrs { f.write_str("; ")?; } - AttrsRep::Map(map) => { + AttrsRep::Im(map) => { for (name, value) in map { write!(f, "{} = ", name.ident_str())?; value.total_fmt(f, set)?; @@ -139,6 +141,7 @@ impl TotalDisplay for NixAttrs { #[cfg(feature = "arbitrary")] mod arbitrary { use super::*; + use std::collections::BTreeMap; use proptest::prelude::*; use proptest::prop_oneof; @@ -158,7 +161,7 @@ mod arbitrary { ) .prop_map(|(name, value)| Self(AttrsRep::KV { name, value })), any_with::>(args) - .prop_map(|map| Self(AttrsRep::Map(map))) + .prop_map(|map| Self::from_iter(map.into_iter())) ] .boxed() } @@ -183,44 +186,41 @@ impl NixAttrs { // Explicitly handle all branches instead of falling // through, to ensure that we get at least some compiler // errors if variants are modified. - (AttrsRep::Map(_), AttrsRep::Map(_)) - | (AttrsRep::Map(_), AttrsRep::KV { .. }) - | (AttrsRep::KV { .. }, AttrsRep::Map(_)) => {} + (AttrsRep::Im(_), AttrsRep::Im(_)) + | (AttrsRep::Im(_), AttrsRep::KV { .. }) + | (AttrsRep::KV { .. }, AttrsRep::Im(_)) => {} }; // Slightly more advanced, but still optimised updates match (self.0, other.0) { - (AttrsRep::Map(mut m), AttrsRep::KV { name, value }) => { + (AttrsRep::Im(mut m), AttrsRep::KV { name, value }) => { m.insert(NixString::NAME, name); m.insert(NixString::VALUE, value); - NixAttrs(AttrsRep::Map(m)) + NixAttrs(AttrsRep::Im(m)) } - (AttrsRep::KV { name, value }, AttrsRep::Map(mut m)) => { + (AttrsRep::KV { name, value }, AttrsRep::Im(mut m)) => { match m.entry(NixString::NAME) { - btree_map::Entry::Vacant(e) => { + imbl::ordmap::Entry::Vacant(e) => { e.insert(name); } - btree_map::Entry::Occupied(_) => { /* name from `m` has precedence */ } + imbl::ordmap::Entry::Occupied(_) => { /* name from `m` has precedence */ } }; match m.entry(NixString::VALUE) { - btree_map::Entry::Vacant(e) => { + imbl::ordmap::Entry::Vacant(e) => { e.insert(value); } - btree_map::Entry::Occupied(_) => { /* value from `m` has precedence */ } + imbl::ordmap::Entry::Occupied(_) => { /* value from `m` has precedence */ } }; - NixAttrs(AttrsRep::Map(m)) + NixAttrs(AttrsRep::Im(m)) } // Plain merge of maps. - (AttrsRep::Map(mut m1), AttrsRep::Map(mut m2)) => { - m1.append(&mut m2); - NixAttrs(AttrsRep::Map(m1)) - } + (AttrsRep::Im(m1), AttrsRep::Im(m2)) => NixAttrs(AttrsRep::Im(m2.union(m1))), // Cases handled above by the borrowing match: _ => unreachable!(), @@ -230,7 +230,7 @@ impl NixAttrs { /// Return the number of key-value entries in an attrset. pub fn len(&self) -> usize { match &self.0 { - AttrsRep::Map(map) => map.len(), + AttrsRep::Im(map) => map.len(), AttrsRep::Empty => 0, AttrsRep::KV { .. } => 2, } @@ -256,7 +256,7 @@ impl NixAttrs { #[allow(clippy::needless_lifetimes)] pub fn iter<'a>(&'a self) -> Iter> { Iter(match &self.0 { - AttrsRep::Map(map) => KeyValue::Map(map.iter()), + AttrsRep::Im(map) => KeyValue::Im(map.iter()), AttrsRep::Empty => KeyValue::Empty, AttrsRep::KV { @@ -280,7 +280,7 @@ impl NixAttrs { ] .into_iter(), )), - AttrsRep::Map(map) => IntoIter(IntoIterRepr::Map(map.into_iter())), + AttrsRep::Im(map) => IntoIter(IntoIterRepr::Im(map.into_iter())), } } @@ -294,7 +294,7 @@ impl NixAttrs { pub fn keys(&self) -> Keys { Keys(match &self.0 { AttrsRep::Empty => KeysInner::Empty, - AttrsRep::Map(m) => KeysInner::Map(m.keys()), + AttrsRep::Im(m) => KeysInner::Im(m.keys()), AttrsRep::KV { .. } => KeysInner::KV(IterKV::default()), }) } @@ -321,7 +321,7 @@ impl NixAttrs { } } - let mut attrs = NixAttrs(AttrsRep::Map(BTreeMap::new())); + let mut attrs = NixAttrs(AttrsRep::Im(OrdMap::new())); for _ in 0..count { let value = stack_slice.pop().unwrap(); @@ -363,7 +363,7 @@ impl NixAttrs { // elements after key construction). In practice this // probably does not happen, so it's better to just bite // the bullet and implement this branch. - (AttrsRep::Empty, AttrsRep::Map(map)) | (AttrsRep::Map(map), AttrsRep::Empty) => { + (AttrsRep::Empty, AttrsRep::Im(map)) | (AttrsRep::Im(map), AttrsRep::Empty) => { Ok(map.is_empty()) } @@ -382,8 +382,8 @@ impl NixAttrs { }, ) => Ok(n1.nix_eq(n2, vm)? && v1.nix_eq(v2, vm)?), - (AttrsRep::Map(map), AttrsRep::KV { name, value }) - | (AttrsRep::KV { name, value }, AttrsRep::Map(map)) => { + (AttrsRep::Im(map), AttrsRep::KV { name, value }) + | (AttrsRep::KV { name, value }, AttrsRep::Im(map)) => { if map.len() != 2 { return Ok(false); } @@ -397,7 +397,7 @@ impl NixAttrs { Ok(false) } - (AttrsRep::Map(m1), AttrsRep::Map(m2)) => { + (AttrsRep::Im(m1), AttrsRep::Im(m2)) => { if m1.len() != m2.len() { return Ok(false); } @@ -462,11 +462,11 @@ fn attempt_optimise_kv(slice: &mut [Value]) -> Option { /// checking against duplicate keys. fn set_attr(attrs: &mut NixAttrs, key: NixString, value: Value) -> Result<(), ErrorKind> { match attrs.0.map_mut().entry(key) { - btree_map::Entry::Occupied(entry) => Err(ErrorKind::DuplicateAttrsKey { + imbl::ordmap::Entry::Occupied(entry) => Err(ErrorKind::DuplicateAttrsKey { key: entry.key().as_str().to_string(), }), - btree_map::Entry::Vacant(entry) => { + imbl::ordmap::Entry::Vacant(entry) => { entry.insert(value); Ok(()) } @@ -495,7 +495,6 @@ impl IterKV { /// Iterator representation over the keys *and* values of an attribute /// set. -#[derive(Debug)] pub enum KeyValue<'a> { Empty, @@ -505,7 +504,7 @@ pub enum KeyValue<'a> { at: IterKV, }, - Map(btree_map::Iter<'a, NixString, Value>), + Im(imbl::ordmap::Iter<'a, NixString, Value>), } /// Iterator over a Nix attribute set. @@ -519,7 +518,7 @@ impl<'a> Iterator for Iter> { fn next(&mut self) -> Option { match &mut self.0 { - KeyValue::Map(inner) => inner.next(), + KeyValue::Im(inner) => inner.next(), KeyValue::Empty => None, KeyValue::KV { name, value, at } => match at { @@ -544,7 +543,7 @@ impl<'a> ExactSizeIterator for Iter> { match &self.0 { KeyValue::Empty => 0, KeyValue::KV { .. } => 2, - KeyValue::Map(inner) => inner.len(), + KeyValue::Im(inner) => inner.len(), } } } @@ -552,7 +551,7 @@ impl<'a> ExactSizeIterator for Iter> { enum KeysInner<'a> { Empty, KV(IterKV), - Map(btree_map::Keys<'a, NixString, Value>), + Im(imbl::ordmap::Keys<'a, NixString, Value>), } pub struct Keys<'a>(KeysInner<'a>); @@ -572,7 +571,7 @@ impl<'a> Iterator for Keys<'a> { Some(NixString::VALUE_REF) } KeysInner::KV(IterKV::Done) => None, - KeysInner::Map(m) => m.next(), + KeysInner::Im(m) => m.next(), } } } @@ -592,7 +591,7 @@ impl<'a> ExactSizeIterator for Keys<'a> { match &self.0 { KeysInner::Empty => 0, KeysInner::KV(_) => 2, - KeysInner::Map(m) => m.len(), + KeysInner::Im(m) => m.len(), } } } @@ -601,7 +600,7 @@ impl<'a> ExactSizeIterator for Keys<'a> { pub enum IntoIterRepr { Empty, Finite(std::vec::IntoIter<(NixString, Value)>), - Map(std::collections::btree_map::IntoIter), + Im(imbl::ordmap::ConsumingIter<(NixString, Value)>), } #[repr(transparent)] @@ -613,8 +612,8 @@ impl Iterator for IntoIter { fn next(&mut self) -> Option { match &mut self.0 { IntoIterRepr::Empty => None, - IntoIterRepr::Map(inner) => inner.next(), IntoIterRepr::Finite(inner) => inner.next(), + IntoIterRepr::Im(inner) => inner.next(), } } } @@ -623,8 +622,8 @@ impl ExactSizeIterator for IntoIter { fn len(&self) -> usize { match &self.0 { IntoIterRepr::Empty => 0, - IntoIterRepr::Map(inner) => inner.len(), IntoIterRepr::Finite(inner) => inner.len(), + IntoIterRepr::Im(inner) => inner.len(), } } } -- cgit 1.4.1