Skip to content

Commit

Permalink
Fix stream realiasing possibly shadowing other streams.
Browse files Browse the repository at this point in the history
  • Loading branch information
adri326 committed Feb 6, 2025
1 parent 949d316 commit 7f2ce57
Show file tree
Hide file tree
Showing 2 changed files with 64 additions and 2 deletions.
19 changes: 19 additions & 0 deletions src/machine/machine_indices.rs
Original file line number Diff line number Diff line change
Expand Up @@ -484,6 +484,8 @@ impl IndexStore {
pub(crate) fn remove_stream(&mut self, stream: Stream) {
if let Some(alias) = stream.options().get_alias() {
debug_assert_eq!(self.stream_aliases.get(&alias), Some(&stream));
assert!(!is_protected_alias(alias));

self.stream_aliases.swap_remove(&alias);
}
self.streams.remove(&stream);
Expand All @@ -503,6 +505,18 @@ impl IndexStore {
callback(options);

if options.get_alias() != prev_alias {
if prev_alias.map(is_protected_alias).unwrap_or(false)
|| options
.get_alias()
.map(|alias| self.has_stream(alias))
.unwrap_or(false)
{
// user_input, user_output and user_error cannot be realiased,
// and realiasing cannot shadow an existing stream.
options.set_alias_to_atom_opt(prev_alias);
return;
}

if let Some(prev_alias) = prev_alias {
self.stream_aliases.swap_remove(&prev_alias);
}
Expand All @@ -516,6 +530,11 @@ impl IndexStore {
self.stream_aliases.contains_key(&alias)
}

/// ## Warning
///
/// The returned stream's options should only be modified through
/// [`IndexStore::update_stream_options`], to avoid breaking the
/// invariants of [`IndexStore`].
pub(crate) fn get_stream(&self, alias: Atom) -> Option<Stream> {
self.stream_aliases.get(&alias).copied()
}
Expand Down
47 changes: 45 additions & 2 deletions src/machine/streams.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1476,6 +1476,10 @@ impl MachineState {
}
}

/// ## Warning
///
/// The options of streams stored in `Machine::indices` should only
/// be modified through [`IndexStore::update_stream_options`].
pub(crate) fn get_stream_options(
&mut self,
alias: HeapCellValue,
Expand Down Expand Up @@ -1591,6 +1595,19 @@ impl MachineState {
options
}

/// If `addr` is a [`Cons`](HeapCellValueTag::Cons) to a stream, then returns it.
///
/// If it is an atom or a string, then this searches for the corresponding stream
/// inside of [`self.indices`], returning it.
///
/// ## Warning
///
/// **Do not directly modify [`stream.options_mut()`](Stream::options_mut)
/// on the returned stream.**
///
/// Other functions rely on the invariants of [`IndexStore`], which may
/// become invalidated by the direct modification of a stream's option (namely,
/// its alias name). Instead, use [`IndexStore::update_stream_options`].
pub(crate) fn get_stream_or_alias(
&mut self,
addr: HeapCellValue,
Expand Down Expand Up @@ -1953,14 +1970,40 @@ mod test {
let mut machine = MachineBuilder::new().build();

let results = machine
.run_query(r#"
.run_query(
r#"
\+ \+ (
open("README.md", read, S, [alias(readme)]),
open(stream(S), read, _, [alias(another_alias)]),
close(S)
),
open("README.md", read, _, [alias(readme)]).
"#)
"#,
)
.collect::<Vec<_>>();

assert_eq!(results.len(), 1);
assert!(results[0].is_ok());
}

#[test]
#[cfg_attr(miri, ignore)]
fn close_realiased_user_output() {
let mut machine = MachineBuilder::new()
.with_streams(StreamConfig::in_memory())
.build();

let results = machine
.run_query(
r#"
\+ \+ (
open("README.md", read, S),
open(stream(S), read, _, [alias(user_output)]),
close(S)
),
write(user_output, hello).
"#,
)
.collect::<Vec<_>>();

assert_eq!(results.len(), 1);
Expand Down

0 comments on commit 7f2ce57

Please sign in to comment.