fix behavior of loops with external stream results (should terminate)

This commit is contained in:
Devyn Cairns 2024-07-08 18:50:56 -07:00
parent 5e04ab2f14
commit fceaefcfaf
4 changed files with 231 additions and 12 deletions

View File

@ -12,6 +12,7 @@ pub(crate) struct BlockBuilder {
pub(crate) ast: Vec<Option<IrAstRef>>,
pub(crate) register_allocation_state: Vec<bool>,
pub(crate) file_count: u32,
pub(crate) loop_stack: Vec<LoopState>,
}
impl BlockBuilder {
@ -24,6 +25,7 @@ impl BlockBuilder {
ast: vec![],
register_allocation_state: vec![true],
file_count: 0,
loop_stack: vec![],
}
}
@ -420,6 +422,69 @@ impl BlockBuilder {
Ok(next)
}
/// Push a new loop state onto the builder.
pub(crate) fn begin_loop(&mut self) {
self.loop_stack.push(LoopState::new());
}
/// True if we are currently in a loop.
pub(crate) fn is_in_loop(&self) -> bool {
!self.loop_stack.is_empty()
}
/// Add a loop breaking jump instruction.
pub(crate) fn push_break(&mut self, span: Span) -> Result<usize, CompileError> {
let index = self.jump_placeholder(span)?;
self.loop_stack
.last_mut()
.ok_or_else(|| CompileError::NotInALoop {
msg: "`break` called from outside of a loop".into(),
span: Some(span),
})?
.break_branches
.push(index);
Ok(index)
}
/// Add a loop continuing jump instruction.
pub(crate) fn push_continue(&mut self, span: Span) -> Result<usize, CompileError> {
let index = self.jump_placeholder(span)?;
self.loop_stack
.last_mut()
.ok_or_else(|| CompileError::NotInALoop {
msg: "`continue` called from outside of a loop".into(),
span: Some(span),
})?
.continue_branches
.push(index);
Ok(index)
}
/// Pop the loop state and set any `break` or `continue` instructions to their appropriate
/// target instruction indexes.
pub(crate) fn end_loop(
&mut self,
break_target_index: usize,
continue_target_index: usize,
) -> Result<(), CompileError> {
let loop_state = self
.loop_stack
.pop()
.ok_or_else(|| CompileError::NotInALoop {
msg: "end_loop() called outside of a loop".into(),
span: None,
})?;
for break_index in loop_state.break_branches {
self.set_branch_target(break_index, break_target_index)?;
}
for continue_index in loop_state.continue_branches {
self.set_branch_target(continue_index, continue_target_index)?;
}
Ok(())
}
/// Consume the builder and produce the final [`IrBlock`].
pub(crate) fn finish(self) -> IrBlock {
IrBlock {
@ -436,3 +501,19 @@ impl BlockBuilder {
}
}
}
/// Keeps track of `break` and `continue` branches that need to be set up after a loop is compiled.
#[derive(Debug)]
pub(crate) struct LoopState {
break_branches: Vec<usize>,
continue_branches: Vec<usize>,
}
impl LoopState {
pub(crate) const fn new() -> Self {
LoopState {
break_branches: vec![],
continue_branches: vec![],
}
}
}

View File

@ -62,6 +62,12 @@ pub(crate) fn compile_call(
"for" => {
return compile_for(working_set, builder, call, redirect_modes, io_reg);
}
"break" => {
return compile_break(working_set, builder, call, redirect_modes, io_reg);
}
"continue" => {
return compile_continue(working_set, builder, call, redirect_modes, io_reg);
}
"return" => {
return compile_return(working_set, builder, call, redirect_modes, io_reg);
}

View File

@ -539,9 +539,13 @@ pub(crate) fn compile_loop(
) -> Result<(), CompileError> {
// Pseudocode:
//
// LOOP: drain %io_reg
// ...<block>...
// drop %io_reg
// LOOP: ...<block>...
// check-external-failed %failed_reg, %io_reg
// drain %io_reg
// branch-if %failed_reg, END
// jump %LOOP
// END: drop %io_reg
let invalid = || CompileError::InvalidKeywordCall {
keyword: "loop".into(),
span: call.head,
@ -551,7 +555,10 @@ pub(crate) fn compile_loop(
let block_id = block_arg.as_block().ok_or_else(invalid)?;
let block = working_set.get_block(block_id);
let loop_index = builder.drain(io_reg, call.head)?;
builder.begin_loop();
builder.load_empty(io_reg)?;
let loop_index = builder.next_instruction_index();
compile_block(
working_set,
@ -562,8 +569,27 @@ pub(crate) fn compile_loop(
io_reg,
)?;
// Check for failed externals, and drain the output
let failed_reg = builder.next_register()?;
builder.push(
Instruction::CheckExternalFailed {
dst: failed_reg,
src: io_reg,
}
.into_spanned(call.head),
)?;
builder.drain(io_reg, call.head)?;
let failed_index = builder.branch_if_placeholder(failed_reg, call.head)?;
builder.jump(loop_index, call.head)?;
let end_index = builder.next_instruction_index();
builder.set_branch_target(failed_index, end_index)?;
builder.end_loop(end_index, loop_index)?;
// State of %io_reg is not necessarily well defined here due to control flow, so make sure it's
// empty.
builder.mark_register(io_reg)?;
builder.load_empty(io_reg)?;
Ok(())
@ -579,13 +605,16 @@ pub(crate) fn compile_while(
) -> Result<(), CompileError> {
// Pseudocode:
//
// LOOP: drain %io_reg
// %io_reg <- <condition>
// drop %io_reg
// LOOP: %io_reg <- <condition>
// branch-if %io_reg, TRUE
// jump FALSE
// TRUE: ...<block>...
// check-external-failed %failed_reg, %io_reg
// drain %io_reg
// branch-if %failed_reg, FALSE
// jump LOOP
// FALSE:
// FALSE: drop %io_reg
let invalid = || CompileError::InvalidKeywordCall {
keyword: "while".into(),
span: call.head,
@ -596,7 +625,10 @@ pub(crate) fn compile_while(
let block_id = block_arg.as_block().ok_or_else(invalid)?;
let block = working_set.get_block(block_id);
let loop_index = builder.drain(io_reg, call.head)?;
builder.begin_loop();
builder.load_empty(io_reg)?;
let loop_index = builder.next_instruction_index();
compile_expression(
working_set,
@ -621,10 +653,29 @@ pub(crate) fn compile_while(
io_reg,
)?;
// We have to check the result for a failed external, and drain the output on each iteration
let failed_reg = builder.next_register()?;
builder.push(
Instruction::CheckExternalFailed {
dst: failed_reg,
src: io_reg,
}
.into_spanned(call.head),
)?;
builder.drain(io_reg, call.head)?;
let failed_index = builder.branch_if_placeholder(failed_reg, call.head)?;
builder.jump(loop_index, call.head)?;
builder.set_branch_target(jump_false_index, builder.next_instruction_index())?;
let end_index = builder.next_instruction_index();
builder.set_branch_target(jump_false_index, end_index)?;
builder.set_branch_target(failed_index, end_index)?;
builder.end_loop(end_index, loop_index)?;
// State of %io_reg is not necessarily well defined here due to control flow, so make sure it's
// empty.
builder.mark_register(io_reg)?;
builder.load_empty(io_reg)?;
Ok(())
@ -644,9 +695,11 @@ pub(crate) fn compile_for(
// LOOP: iterate %io_reg, %stream_reg, END
// store-variable $var, %io_reg
// %io_reg <- <...block...>
// check-external-failed %failed_reg, %io_reg
// drain %io_reg
// branch-if %failed_reg, END
// jump LOOP
// END:
// END: drop %io_reg
let invalid = || CompileError::InvalidKeywordCall {
keyword: "for".into(),
span: call.head,
@ -670,6 +723,9 @@ pub(crate) fn compile_for(
// Ensure io_reg is marked so we don't use it
builder.mark_register(io_reg)?;
// Set up loop state
builder.begin_loop();
let stream_reg = builder.next_register()?;
compile_expression(
@ -711,24 +767,90 @@ pub(crate) fn compile_for(
io_reg,
)?;
// Check for failed external
let failed_reg = builder.next_register()?;
builder.push(
Instruction::CheckExternalFailed {
dst: failed_reg,
src: io_reg,
}
.into_spanned(call.head),
)?;
// Drain the output
builder.drain(io_reg, call.head)?;
// End the loop if there was a failure
let failed_index = builder.branch_if_placeholder(failed_reg, call.head)?;
// Loop back to iterate to get the next value
builder.jump(iterate_index, call.head)?;
// Update the iterate target to the end of the loop
let target_index = builder.next_instruction_index();
builder.set_branch_target(iterate_index, target_index)?;
let end_index = builder.next_instruction_index();
builder.set_branch_target(failed_index, end_index)?;
builder.set_branch_target(iterate_index, end_index)?;
// Handle break/continue
builder.end_loop(end_index, iterate_index)?;
// We don't need stream_reg anymore, after the loop
// io_reg is guaranteed to be Empty due to the iterate instruction before
// io_reg may or may not be empty, so be sure it is
builder.free_register(stream_reg)?;
builder.mark_register(io_reg)?;
builder.load_empty(io_reg)?;
Ok(())
}
/// Compile a call to `break`.
pub(crate) fn compile_break(
_working_set: &StateWorkingSet,
builder: &mut BlockBuilder,
call: &Call,
_redirect_modes: RedirectModes,
io_reg: RegId,
) -> Result<(), CompileError> {
if builder.is_in_loop() {
builder.load_empty(io_reg)?;
builder.push_break(call.head)?;
} else {
// Fall back to calling the command if we can't find the loop target statically
builder.push(
Instruction::Call {
decl_id: call.decl_id,
src_dst: io_reg,
}
.into_spanned(call.head),
)?;
}
Ok(())
}
/// Compile a call to `continue`.
pub(crate) fn compile_continue(
_working_set: &StateWorkingSet,
builder: &mut BlockBuilder,
call: &Call,
_redirect_modes: RedirectModes,
io_reg: RegId,
) -> Result<(), CompileError> {
if builder.is_in_loop() {
builder.load_empty(io_reg)?;
builder.push_continue(call.head)?;
} else {
// Fall back to calling the command if we can't find the loop target statically
builder.push(
Instruction::Call {
decl_id: call.decl_id,
src_dst: io_reg,
}
.into_spanned(call.head),
)?;
}
Ok(())
}
/// Compile a call to `return` as a `return` instruction.
///
/// This is not strictly necessary, but it is more efficient.

View File

@ -136,6 +136,7 @@ pub enum CompileError {
},
#[error("Missing required declaration: `{decl_name}`")]
#[diagnostic(code(nu::compile::missing_required_declaration))]
MissingRequiredDeclaration {
decl_name: String,
#[label("`{decl_name}` must be in scope to compile this expression")]
@ -143,9 +144,18 @@ pub enum CompileError {
},
#[error("Invalid literal")]
#[diagnostic(code(nu::compile::invalid_literal))]
InvalidLiteral {
msg: String,
#[label("{msg}")]
span: Span,
},
#[error("{msg}")]
#[diagnostic(code(nu::compile::not_in_a_loop))]
NotInALoop {
msg: String,
#[label("can't be used outside of a loop")]
span: Option<Span>,
},
}