feat: add exit operation to prelude

ss2165 · ss2165 · commit 5e3cbace7941 · 2025-03-20T12:56:08.000Z
panic ends all shots, exit ends current shot and signals next shot can still run

By default they are handled identically with an abort in codegen - I am not sure if this is the right thing to do.
diff --git a/hugr-core/src/extension/prelude.rs b/hugr-core/src/extension/prelude.rs
@@ -102,13 +102,27 @@ lazy_static! {
                     PolyFuncTypeRV::new(
                         [TypeParam::new_list(TypeBound::Any), TypeParam::new_list(TypeBound::Any)],
                         FuncValueType::new(
-                            vec![TypeRV::new_extension(error_type), TypeRV::new_row_var_use(0, TypeBound::Any)],
+                            vec![TypeRV::new_extension(error_type.clone()), TypeRV::new_row_var_use(0, TypeBound::Any)],
                             vec![TypeRV::new_row_var_use(1, TypeBound::Any)],
                         ),
                     ),
                     extension_ref,
                 )
                 .unwrap();
+            prelude
+            .add_op(
+                EXIT_OP_ID,
+                "Exit with input error".to_string(),
+                PolyFuncTypeRV::new(
+                    [TypeParam::new_list(TypeBound::Any), TypeParam::new_list(TypeBound::Any)],
+                    FuncValueType::new(
+                        vec![TypeRV::new_extension(error_type), TypeRV::new_row_var_use(0, TypeBound::Any)],
+                        vec![TypeRV::new_row_var_use(1, TypeBound::Any)],
+                    ),
+                ),
+                extension_ref,
+            )
+            .unwrap();
 
             TupleOpDef::load_all_ops(prelude, extension_ref).unwrap();
             NoopDef.add_to_extension(prelude, extension_ref).unwrap();
@@ -163,8 +177,25 @@ pub fn bool_t() -> Type {
 /// second sequence of types in its instantiation. Note that the inputs and
 /// outputs only exist so that structural constraints such as linearity can be
 /// satisfied.
+///
+/// Panic immediately halts a multi-shot program. It is intended to be used in
+/// cases where an unrecoverable error is detected.
 pub const PANIC_OP_ID: OpName = OpName::new_inline("panic");
 
+/// Name of the prelude exit operation.
+///
+/// This operation can have any input and any output wires; it is instantiated
+/// with two [TypeArg::Sequence]s representing these. The first input to the
+/// operation is always an error type; the remaining inputs correspond to the
+/// first sequence of types in its instantiation; the outputs correspond to the
+/// second sequence of types in its instantiation. Note that the inputs and
+/// outputs only exist so that structural constraints such as linearity can be
+/// satisfied.
+///
+/// Exit immediately halts a single shot's execution. It is intended to be used
+/// when the user wants to exit the program early.
+pub const EXIT_OP_ID: OpName = OpName::new_inline("exit");
+
 /// Name of the string type.
 pub const STRING_TYPE_NAME: TypeName = TypeName::new_inline("string");
 
@@ -1077,7 +1108,7 @@ mod test {
 
         const TYPE_ARG_NONE: TypeArg = TypeArg::Sequence { elems: vec![] };
         let op = PRELUDE
-            .instantiate_extension_op(&PANIC_OP_ID, [TYPE_ARG_NONE, TYPE_ARG_NONE])
+            .instantiate_extension_op(&EXIT_OP_ID, [TYPE_ARG_NONE, TYPE_ARG_NONE])
             .unwrap();
 
         b.add_dataflow_op(op, [err]).unwrap();
diff --git a/hugr-llvm/src/extension/prelude.rs b/hugr-llvm/src/extension/prelude.rs
@@ -340,6 +340,26 @@ pub fn add_prelude_extensions<'a, H: HugrView<Node = Node> + 'a>(
             args.outputs.finish(context.builder(), returns)
         }
     })
+    .extension_op(prelude::PRELUDE_ID, prelude::EXIT_OP_ID, {
+        // by default treat an exit like a panic
+        let pcg = pcg.clone();
+        move |context, args| {
+            let err = args.inputs[0];
+            ensure!(
+                err.get_type()
+                    == pcg
+                        .error_type(&context.typing_session())?
+                        .as_basic_type_enum()
+            );
+            pcg.emit_panic(context, err)?;
+            let returns = args
+                .outputs
+                .get_types()
+                .map(|ty| ty.const_zero())
+                .collect_vec();
+            args.outputs.finish(context.builder(), returns)
+        }
+    })
     .extension_op(prelude::PRELUDE_ID, generic::LOAD_NAT_OP_ID, {
         let pcg = pcg.clone();
         move |context, args| {
@@ -362,6 +382,7 @@ pub fn add_prelude_extensions<'a, H: HugrView<Node = Node> + 'a>(
 #[cfg(test)]
 mod test {
     use hugr_core::builder::{Dataflow, DataflowSubContainer};
+    use hugr_core::extension::prelude::EXIT_OP_ID;
     use hugr_core::extension::PRELUDE;
     use hugr_core::types::{Type, TypeArg};
     use hugr_core::{type_row, Hugr};
@@ -522,6 +543,34 @@ mod test {
         check_emission!(hugr, prelude_llvm_ctx);
     }
 
+    #[rstest]
+    fn prelude_exit(prelude_llvm_ctx: TestContext) {
+        let error_val = ConstError::new(42, "EXIT");
+        let type_arg_q: TypeArg = TypeArg::Type { ty: qb_t() };
+        let type_arg_2q: TypeArg = TypeArg::Sequence {
+            elems: vec![type_arg_q.clone(), type_arg_q],
+        };
+        let exit_op = PRELUDE
+            .instantiate_extension_op(&EXIT_OP_ID, [type_arg_2q.clone(), type_arg_2q.clone()])
+            .unwrap();
+
+        let hugr = SimpleHugrConfig::new()
+            .with_ins(vec![qb_t(), qb_t()])
+            .with_outs(vec![qb_t(), qb_t()])
+            .with_extensions(prelude::PRELUDE_REGISTRY.to_owned())
+            .finish(|mut builder| {
+                let [q0, q1] = builder.input_wires_arr();
+                let err = builder.add_load_value(error_val);
+                let [q0, q1] = builder
+                    .add_dataflow_op(exit_op, [err, q0, q1])
+                    .unwrap()
+                    .outputs_arr();
+                builder.finish_with_outputs([q0, q1]).unwrap()
+            });
+
+        check_emission!(hugr, prelude_llvm_ctx);
+    }
+
     #[rstest]
     fn prelude_print(prelude_llvm_ctx: TestContext) {
         let greeting: ConstString = ConstString::new("Hello, world!".into());
diff --git a/hugr-llvm/src/extension/snapshots/hugr_llvm__extension__prelude__test__prelude_exit@llvm14.snap b/hugr-llvm/src/extension/snapshots/hugr_llvm__extension__prelude__test__prelude_exit@llvm14.snap
@@ -0,0 +1,27 @@
+---
+source: hugr-llvm/src/extension/prelude.rs
+expression: mod_str
+---
+; ModuleID = 'test_context'
+source_filename = "test_context"
+
+@0 = private unnamed_addr constant [5 x i8] c"EXIT\00", align 1
+@prelude.panic_template = private unnamed_addr constant [34 x i8] c"Program panicked (signal %i): %s\0A\00", align 1
+
+define { i16, i16 } @_hl.main.1(i16 %0, i16 %1) {
+alloca_block:
+  br label %entry_block
+
+entry_block:                                      ; preds = %alloca_block
+  %2 = extractvalue { i32, i8* } { i32 42, i8* getelementptr inbounds ([5 x i8], [5 x i8]* @0, i32 0, i32 0) }, 0
+  %3 = extractvalue { i32, i8* } { i32 42, i8* getelementptr inbounds ([5 x i8], [5 x i8]* @0, i32 0, i32 0) }, 1
+  %4 = call i32 (i8*, ...) @printf(i8* getelementptr inbounds ([34 x i8], [34 x i8]* @prelude.panic_template, i32 0, i32 0), i32 %2, i8* %3)
+  call void @abort()
+  %mrv = insertvalue { i16, i16 } undef, i16 0, 0
+  %mrv8 = insertvalue { i16, i16 } %mrv, i16 0, 1
+  ret { i16, i16 } %mrv8
+}
+
+declare i32 @printf(i8*, ...)
+
+declare void @abort()
diff --git a/hugr-llvm/src/extension/snapshots/hugr_llvm__extension__prelude__test__prelude_exit@pre-mem2reg@llvm14.snap b/hugr-llvm/src/extension/snapshots/hugr_llvm__extension__prelude__test__prelude_exit@pre-mem2reg@llvm14.snap
@@ -0,0 +1,48 @@
+---
+source: hugr-llvm/src/extension/prelude.rs
+expression: mod_str
+---
+; ModuleID = 'test_context'
+source_filename = "test_context"
+
+@0 = private unnamed_addr constant [5 x i8] c"EXIT\00", align 1
+@prelude.panic_template = private unnamed_addr constant [34 x i8] c"Program panicked (signal %i): %s\0A\00", align 1
+
+define { i16, i16 } @_hl.main.1(i16 %0, i16 %1) {
+alloca_block:
+  %"0" = alloca i16, align 2
+  %"1" = alloca i16, align 2
+  %"5_0" = alloca { i32, i8* }, align 8
+  %"2_0" = alloca i16, align 2
+  %"2_1" = alloca i16, align 2
+  %"6_0" = alloca i16, align 2
+  %"6_1" = alloca i16, align 2
+  br label %entry_block
+
+entry_block:                                      ; preds = %alloca_block
+  store { i32, i8* } { i32 42, i8* getelementptr inbounds ([5 x i8], [5 x i8]* @0, i32 0, i32 0) }, { i32, i8* }* %"5_0", align 8
+  store i16 %0, i16* %"2_0", align 2
+  store i16 %1, i16* %"2_1", align 2
+  %"5_01" = load { i32, i8* }, { i32, i8* }* %"5_0", align 8
+  %"2_02" = load i16, i16* %"2_0", align 2
+  %"2_13" = load i16, i16* %"2_1", align 2
+  %2 = extractvalue { i32, i8* } %"5_01", 0
+  %3 = extractvalue { i32, i8* } %"5_01", 1
+  %4 = call i32 (i8*, ...) @printf(i8* getelementptr inbounds ([34 x i8], [34 x i8]* @prelude.panic_template, i32 0, i32 0), i32 %2, i8* %3)
+  call void @abort()
+  store i16 0, i16* %"6_0", align 2
+  store i16 0, i16* %"6_1", align 2
+  %"6_04" = load i16, i16* %"6_0", align 2
+  %"6_15" = load i16, i16* %"6_1", align 2
+  store i16 %"6_04", i16* %"0", align 2
+  store i16 %"6_15", i16* %"1", align 2
+  %"06" = load i16, i16* %"0", align 2
+  %"17" = load i16, i16* %"1", align 2
+  %mrv = insertvalue { i16, i16 } undef, i16 %"06", 0
+  %mrv8 = insertvalue { i16, i16 } %mrv, i16 %"17", 1
+  ret { i16, i16 } %mrv8
+}
+
+declare i32 @printf(i8*, ...)
+
+declare void @abort()
diff --git a/hugr-py/src/hugr/std/_json_defs/prelude.json b/hugr-py/src/hugr/std/_json_defs/prelude.json
@@ -195,6 +195,54 @@
       },
       "binary": false
     },
+    "exit": {
+      "extension": "prelude",
+      "name": "exit",
+      "description": "Exit with input error",
+      "signature": {
+        "params": [
+          {
+            "tp": "List",
+            "param": {
+              "tp": "Type",
+              "b": "A"
+            }
+          },
+          {
+            "tp": "List",
+            "param": {
+              "tp": "Type",
+              "b": "A"
+            }
+          }
+        ],
+        "body": {
+          "input": [
+            {
+              "t": "Opaque",
+              "extension": "prelude",
+              "id": "error",
+              "args": [],
+              "bound": "C"
+            },
+            {
+              "t": "R",
+              "i": 0,
+              "b": "A"
+            }
+          ],
+          "output": [
+            {
+              "t": "R",
+              "i": 1,
+              "b": "A"
+            }
+          ],
+          "runtime_reqs": []
+        }
+      },
+      "binary": false
+    },
     "load_nat": {
       "extension": "prelude",
       "name": "load_nat",
diff --git a/specification/std_extensions/prelude.json b/specification/std_extensions/prelude.json
@@ -195,6 +195,54 @@
       },
       "binary": false
     },
+    "exit": {
+      "extension": "prelude",
+      "name": "exit",
+      "description": "Exit with input error",
+      "signature": {
+        "params": [
+          {
+            "tp": "List",
+            "param": {
+              "tp": "Type",
+              "b": "A"
+            }
+          },
+          {
+            "tp": "List",
+            "param": {
+              "tp": "Type",
+              "b": "A"
+            }
+          }
+        ],
+        "body": {
+          "input": [
+            {
+              "t": "Opaque",
+              "extension": "prelude",
+              "id": "error",
+              "args": [],
+              "bound": "C"
+            },
+            {
+              "t": "R",
+              "i": 0,
+              "b": "A"
+            }
+          ],
+          "output": [
+            {
+              "t": "R",
+              "i": 1,
+              "b": "A"
+            }
+          ],
+          "runtime_reqs": []
+        }
+      },
+      "binary": false
+    },
     "load_nat": {
       "extension": "prelude",
       "name": "load_nat",
