From 649a862ae14d334832831725cf712f737891032c Mon Sep 17 00:00:00 2001 From: Florian Klink Date: Thu, 23 May 2024 15:10:32 +0200 Subject: feat(tvix/eval): rm NixContext::join, add take_context & IntoIterator In places where we want to extend context with that from another NixString, use take_context() to split it off, then call .extend(), making use of IntoIterator to avoid a bunch of clones. Change-Id: I2460141a3ed776c64c36132b2203b6a1d710b922 Reviewed-on: https://cl.tvl.fyi/c/depot/+/11705 Tested-by: BuildkiteCI Autosubmit: flokli Reviewed-by: edef --- tvix/eval/src/builtins/mod.rs | 58 +++++++++++++++--------------------- tvix/eval/src/value/json.rs | 8 ++--- tvix/eval/src/value/mod.rs | 4 +-- tvix/eval/src/value/string.rs | 39 ++++++++++++++---------- tvix/eval/src/vm/mod.rs | 4 +-- tvix/glue/src/builtins/derivation.rs | 4 +-- 6 files changed, 57 insertions(+), 60 deletions(-) diff --git a/tvix/eval/src/builtins/mod.rs b/tvix/eval/src/builtins/mod.rs index 2178d9c44f..5cd94bcf37 100644 --- a/tvix/eval/src/builtins/mod.rs +++ b/tvix/eval/src/builtins/mod.rs @@ -274,9 +274,10 @@ mod pure_builtins { list: Value, ) -> Result { let mut separator = separator.to_contextful_str()?; + let mut context = NixContext::new(); - if let Some(sep_context) = separator.context_mut() { - context = context.join(sep_context); + if let Some(sep_context) = separator.take_context() { + context.extend(sep_context.into_iter()) } let list = list.to_list()?; let mut res = BString::default(); @@ -296,13 +297,8 @@ mod pure_builtins { { Ok(mut s) => { res.push_str(&s); - if let Some(ref mut other_context) = s.context_mut() { - // It is safe to consume the other context here - // because the `list` and `separator` are originally - // moved, here. - // We are not going to use them again - // because the result here is a string. - context = context.join(other_context); + if let Some(other_context) = s.take_context() { + context.extend(other_context.into_iter()); } } Err(c) => return Ok(Value::Catchable(Box::new(c))), @@ -764,9 +760,8 @@ mod pure_builtins { } if let Some(origin_ctx) = origin.context_mut() { - // FUTUREWORK(performance): avoid this clone - // and extend in-place. - *origin_ctx = origin_ctx.clone().join(&mut ctx_elements.into()); + origin_ctx.extend(ctx_elements) + // TODO: didn't we forget cases where origin had no context? } Ok(origin.into()) @@ -1169,8 +1164,8 @@ mod pure_builtins { let mut empty_string_replace = false; let mut context = NixContext::new(); - if let Some(string_context) = string.context_mut() { - context = context.join(string_context); + if let Some(string_context) = string.take_context() { + context.extend(string_context.into_iter()); } // This can't be implemented using Rust's string.replace() as @@ -1200,8 +1195,8 @@ mod pure_builtins { if string[i..i + from.len()] == *from { res.push_str(&to); i += from.len(); - if let Some(to_ctx) = to.context_mut() { - context = context.join(to_ctx); + if let Some(to_ctx) = to.take_context() { + context.extend(to_ctx.into_iter()); } // remember if we applied the empty from->to @@ -1232,8 +1227,8 @@ mod pure_builtins { if from.is_empty() { res.push_str(&to); - if let Some(to_ctx) = to.context_mut() { - context = context.join(to_ctx); + if let Some(to_ctx) = to.take_context() { + context.extend(to_ctx.into_iter()); } break; } @@ -1640,6 +1635,8 @@ pub fn pure_builtins() -> Vec<(&'static str, Value)> { #[builtins] mod placeholder_builtins { + use crate::NixContext; + use super::*; #[builtin("unsafeDiscardStringContext")] @@ -1688,24 +1685,17 @@ mod placeholder_builtins { .to_contextful_str()?; // If there's any context, we will swap any ... by a path one. - if let Some(ctx) = v.context_mut() { - let new_context: tvix_eval::NixContext = ctx - .iter() - .map(|elem| match elem { - // FUTUREWORK(performance): ideally, we should either: - // (a) do interior mutation of the existing context. - // (b) let the structural sharing make those clones cheap. - crate::NixContextElement::Derivation(drv_path) => { - crate::NixContextElement::Plain(drv_path.to_string()) - } - elem => elem.clone(), - }) - .collect::>() - .into(); + if let Some(c) = v.take_context() { + let mut context = NixContext::new(); + context.extend(c.into_iter().map(|elem| match elem { + crate::NixContextElement::Derivation(drv_path) => { + crate::NixContextElement::Plain(drv_path.to_string()) + } + elem => elem.clone(), + })); - *ctx = new_context; + return Ok(Value::String(NixString::new_context_from(context, v))); } - Ok(Value::from(v)) } diff --git a/tvix/eval/src/value/json.rs b/tvix/eval/src/value/json.rs index c48e9c1f4e..24a6bcaf6f 100644 --- a/tvix/eval/src/value/json.rs +++ b/tvix/eval/src/value/json.rs @@ -47,8 +47,8 @@ impl Value { for val in l.into_iter() { match generators::request_to_json(co, val).await { - Ok((v, mut ctx)) => { - context = context.join(&mut ctx); + Ok((v, ctx)) => { + context.extend(ctx.into_iter()); out.push(v) } Err(cek) => return Ok(Err(cek)), @@ -100,8 +100,8 @@ impl Value { out.insert( name.to_str()?.to_owned(), match generators::request_to_json(co, value).await { - Ok((v, mut ctx)) => { - context = context.join(&mut ctx); + Ok((v, ctx)) => { + context.extend(ctx.into_iter()); v } Err(cek) => return Ok(Err(cek)), diff --git a/tvix/eval/src/value/mod.rs b/tvix/eval/src/value/mod.rs index c171c9a04e..dfad0cd839 100644 --- a/tvix/eval/src/value/mod.rs +++ b/tvix/eval/src/value/mod.rs @@ -338,8 +338,8 @@ impl Value { let coerced: Result = match (value, kind) { // coercions that are always done (Value::String(mut s), _) => { - if let Some(ctx) = s.context_mut() { - context = context.join(ctx); + if let Some(ctx) = s.take_context() { + context.extend(ctx.into_iter()); } Ok((*s).into()) } diff --git a/tvix/eval/src/value/string.rs b/tvix/eval/src/value/string.rs index a418a276f9..2a7a202096 100644 --- a/tvix/eval/src/value/string.rs +++ b/tvix/eval/src/value/string.rs @@ -78,15 +78,6 @@ impl NixContext { self } - /// Consumes both ends of the join into a new NixContent - /// containing the union of elements of both ends. - pub fn join(mut self, other: &mut NixContext) -> Self { - let other_set = std::mem::take(&mut other.0); - let mut set: HashSet = std::mem::take(&mut self.0); - set.extend(other_set); - Self(set) - } - /// Extends the existing context with more context elements. pub fn extend(&mut self, iter: T) where @@ -162,6 +153,16 @@ impl NixContext { } } +impl IntoIterator for NixContext { + type Item = NixContextElement; + + type IntoIter = std::collections::hash_set::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.0.into_iter() + } +} + /// This type is never instantiated, but serves to document the memory layout of the actual heap /// allocation for Nix strings. #[allow(dead_code)] @@ -715,8 +716,10 @@ impl NixString { let context = [self.context(), other.context()] .into_iter() .flatten() - .fold(NixContext::new(), |acc_ctx, new_ctx| { - acc_ctx.join(&mut new_ctx.clone()) + .fold(NixContext::new(), |mut acc_ctx, new_ctx| { + // TODO: consume new_ctx? + acc_ctx.extend(new_ctx.iter().cloned()); + acc_ctx }); Self::new_context_from(context, s) } @@ -730,13 +733,13 @@ impl NixString { unsafe { NixStringInner::context_ref(self.0).as_deref() } } - pub(crate) fn context_mut(&mut self) -> Option<&mut NixContext> { + pub(crate) fn context_mut(&mut self) -> &mut Option> { // SAFETY: There's no way to construct an uninitialized or invalid NixString (see the SAFETY // comment in `new`). // // Also, we're using the same lifetime and mutability as self, to fit the // pointer-to-reference conversion rules - unsafe { NixStringInner::context_mut(self.0).as_deref_mut() } + unsafe { NixStringInner::context_mut(self.0) } } pub fn iter_context(&self) -> impl Iterator { @@ -764,12 +767,16 @@ impl NixString { self.context().is_some() } + /// This clears the context of the string, returning + /// the removed dependency tracking information. + pub fn take_context(&mut self) -> Option> { + self.context_mut().take() + } + /// This clears the context of that string, losing /// all dependency tracking information. pub fn clear_context(&mut self) { - // SAFETY: There's no way to construct an uninitialized or invalid NixString (see the SAFETY - // comment in `new`). - *unsafe { NixStringInner::context_mut(self.0) } = None; + let _ = self.take_context(); } pub fn chars(&self) -> Chars<'_> { diff --git a/tvix/eval/src/vm/mod.rs b/tvix/eval/src/vm/mod.rs index 5c244cc3ca..1fa9eba4bc 100644 --- a/tvix/eval/src/vm/mod.rs +++ b/tvix/eval/src/vm/mod.rs @@ -994,8 +994,8 @@ where } let mut nix_string = val.to_contextful_str().with_span(frame, self)?; out.push_str(nix_string.as_bstr()); - if let Some(nix_string_ctx) = nix_string.context_mut() { - context = context.join(nix_string_ctx); + if let Some(nix_string_ctx) = nix_string.take_context() { + context.extend(nix_string_ctx.into_iter()) } } diff --git a/tvix/glue/src/builtins/derivation.rs b/tvix/glue/src/builtins/derivation.rs index a7742ae40a..7878ce0ed9 100644 --- a/tvix/glue/src/builtins/derivation.rs +++ b/tvix/glue/src/builtins/derivation.rs @@ -372,12 +372,12 @@ pub(crate) mod derivation_builtins { return Ok(val); } - let (val_json, mut context) = match val.into_contextful_json(&co).await? { + let (val_json, context) = match val.into_contextful_json(&co).await? { Ok(v) => v, Err(cek) => return Ok(Value::from(cek)), }; - input_context = input_context.join(&mut context); + input_context.extend(context.into_iter()); // No need to check for dups, we only iterate over every attribute name once structured_attrs.insert(arg_name.to_owned(), val_json); -- cgit 1.4.1