Heap memory limit

build/instructions_template.rs

@@ -607,6 +607,10 @@ enum SystemClauseType {
     #[strum_discriminants(strum(props(Arity = "1", Name = "$argv")))]
     Argv,
     REPL(REPLCodePtr),
+    #[strum_discriminants(strum(props(Arity = "1", Name = "$set_heap_limit")))]
+    SetHeapLimit,
+    #[strum_discriminants(strum(props(Arity = "1", Name = "$get_heap_limit")))]
+    GetHeapLimit,
 }
 
 #[allow(dead_code)]
@@ -1872,6 +1876,8 @@ fn generate_instruction_preface() -> TokenStream {
                     &Instruction::CallAddNonCountedBacktracking |
                     &Instruction::CallPopCount |
                     &Instruction::CallArgv |
+                    &Instruction::CallSetHeapLimit |
+                    &Instruction::CallGetHeapLimit |
                     &Instruction::CallEd25519SignRaw |
                     &Instruction::CallEd25519VerifyRaw |
                     &Instruction::CallEd25519SeedToPublicKey => {
@@ -2108,6 +2114,8 @@ fn generate_instruction_preface() -> TokenStream {
                     &Instruction::ExecuteAddNonCountedBacktracking |
                     &Instruction::ExecutePopCount |
                     &Instruction::ExecuteArgv |
+                    &Instruction::ExecuteSetHeapLimit |
+                    &Instruction::ExecuteGetHeapLimit |
                     &Instruction::ExecuteEd25519SignRaw |
                     &Instruction::ExecuteEd25519VerifyRaw |
                     &Instruction::ExecuteEd25519SeedToPublicKey => {

src/machine/dispatch.rs

@@ -4155,6 +4155,22 @@ impl Machine {
                         try_or_throw!(self.machine_st, self.argv());
                         step_or_fail!(self, self.machine_st.p = self.machine_st.cp);
                     }
+                    &Instruction::CallSetHeapLimit => {
+                        try_or_throw!(self.machine_st, self.set_heap_limit());
+                        step_or_fail!(self, self.machine_st.p += 1);
+                    }
+                    &Instruction::ExecuteSetHeapLimit => {
+                        try_or_throw!(self.machine_st, self.set_heap_limit());
+                        step_or_fail!(self, self.machine_st.p = self.machine_st.cp);
+                    }
+                    &Instruction::CallGetHeapLimit => {
+                        try_or_throw!(self.machine_st, self.get_heap_limit());
+                        step_or_fail!(self, self.machine_st.p += 1);
+                    }
+                    &Instruction::ExecuteGetHeapLimit => {
+                        try_or_throw!(self.machine_st, self.get_heap_limit());
+                        step_or_fail!(self, self.machine_st.p = self.machine_st.cp);
+                    }
                     &Instruction::CallCurrentTime => {
                         self.current_time();
                         step_or_fail!(self, self.machine_st.p += 1);
@@ -5202,6 +5218,11 @@ impl Machine {
                 }
             }
 
+            // Not sure if this is the best place to put this
+            if let Err(()) = self.machine_st.check_heap_limit() {
+                continue;
+            }
+
             let interrupted = INTERRUPT.load(std::sync::atomic::Ordering::Relaxed);
 
             match INTERRUPT.compare_exchange(

src/machine/machine_errors.rs

@@ -79,6 +79,7 @@ impl ValidType {
 
 #[derive(Debug, Clone, Copy)]
 pub(crate) enum ResourceError {
+    HeapLimit,
     FiniteMemory(HeapCellValue),
     OutOfFiles,
 }
@@ -325,6 +326,15 @@ impl MachineState {
             ResourceError::OutOfFiles => {
                 functor!(atom!("resource_error"), [atom(atom!("file_descriptors"))])
             }
+            ResourceError::HeapLimit => {
+                functor!(
+                    atom!("resource_error"),
+                    [
+                        atom(atom!("heap_limit")),
+                        fixnum(self.heap_limit.expect("should have heap limit"))
+                    ]
+                )
+            }
         };
 
         MachineError {
@@ -667,7 +677,7 @@ impl MachineState {
 }
 
 impl MachineError {
-    fn into_iter(self, offset: usize) -> Box<dyn Iterator<Item = HeapCellValue>> {
+    pub(super) fn into_iter(self, offset: usize) -> Box<dyn Iterator<Item = HeapCellValue>> {
         match self.from {
             ErrorProvenance::Constructed => {
                 Box::new(self.stub.into_iter().map(move |hcv| hcv + offset))

src/machine/machine_state.rs

@@ -74,6 +74,7 @@ pub struct MachineState {
     pub(super) attr_var_init: AttrVarInitializer,
     pub(super) fail: bool,
     pub heap: Heap,
+    pub(super) heap_limit: Option<usize>,
     pub(super) mode: MachineMode,
     pub(crate) stack: Stack,
     pub(super) registers: Registers,

src/machine/machine_state_impl.rs

@@ -39,6 +39,7 @@ impl MachineState {
             attr_var_init: AttrVarInitializer::new(0),
             fail: false,
             heap: Heap::with_capacity(256 * 256),
+            heap_limit: None,
             mode: MachineMode::Write,
             stack: Stack::new(),
             registers: [heap_loc_as_cell!(0); MAX_ARITY + 1], // self.registers[0] is never used.
@@ -1723,4 +1724,28 @@ impl MachineState {
 
         self.throw_exception(err);
     }
+
+    pub fn check_heap_limit(&mut self) -> Result<(), ()> {
+        if let Some(heap_limit) = self.heap_limit {
+            let h = self.heap.len();
+            if h > heap_limit / 8 {
+                let error = self.resource_error(ResourceError::HeapLimit);
+
+                let mut stub = vec![
+                    atom_as_cell!(atom!("error"), 2),
+                    str_loc_as_cell!(h + 3),
+                    heap_loc_as_cell!(h + 2),
+                ];
+                stub.extend(error.into_iter(3));
+
+                self.throw_exception(stub);
+                self.backtrack();
+                Err(())
+            } else {
+                Ok(())
+            }
+        } else {
+            Err(())
+        }
+    }
 }

src/machine/system_calls.rs

@@ -4971,6 +4971,37 @@ impl Machine {
         Ok(())
     }
 
+    #[inline(always)]
+    pub(crate) fn set_heap_limit(&mut self) -> CallResult {
+        let value = self.deref_register(1);
+
+        let stub = functor_stub(atom!("$set_heap_limit"), 1);
+        let error = stub.type_error(&mut self.machine_st, ValidType::Integer);
+
+        value
+            .to_fixnum()
+            .ok_or(self.machine_st.error_form(error, stub))
+            .map(|fixnum| {
+                self.machine_st.heap_limit = Some(fixnum.get_num() as usize);
+            })
+    }
+
+    #[inline(always)]
+    pub(crate) fn get_heap_limit(&mut self) -> CallResult {
+        let value = self.deref_register(1);
+
+        let cell = self
+            .machine_st
+            .heap_limit
+            .map(|limit| fixnum_as_cell!(Fixnum::build_with(limit as i64)))
+            .unwrap_or(atom_as_cell!(atom!("no_limit")));
+        self.machine_st.heap.push(cell);
+
+        let cell_loc = heap_loc_as_cell!(self.machine_st.heap.len() - 1);
+        unify!(self.machine_st, value, cell_loc);
+        Ok(())
+    }
+
     #[inline(always)]
     pub(crate) fn current_time(&mut self) {
         let timestamp = self.systemtime_to_timestamp(SystemTime::now());