add max arity checks at various stages (#1459)

src/allocator.rs

@@ -59,6 +59,7 @@ pub(crate) trait Allocator {
     fn bindings_mut(&mut self) -> &mut AllocVarDict;
 
     fn take_bindings(self) -> AllocVarDict;
+    fn max_reg_allocated(&self) -> usize;
 
     fn drain_var_data<'a>(
         &mut self,

src/codegen.rs

@@ -865,6 +865,10 @@ impl<'b> CodeGenerator<'b> {
                         };
 
                         self.compile_query_line(term, term_loc, code, num_perm_vars, is_exposed);
+
+                        if self.marker.max_reg_allocated() > MAX_ARITY {
+                            return Err(CompilationError::from(ParserError::ExceededMaxArity));
+                        }
                     }
                 }
             }
@@ -934,6 +938,10 @@ impl<'b> CodeGenerator<'b> {
         let iter = FactIterator::from_rule_head_clause(args);
         let mut fact = self.compile_target::<FactInstruction, _>(iter, GenContext::Head, false);
 
+        if self.marker.max_reg_allocated() > MAX_ARITY {
+            return Err(CompilationError::from(ParserError::ExceededMaxArity));
+        }
+
         let mut unsafe_var_marker = UnsafeVarMarker::new();
 
         if !fact.is_empty() {
@@ -971,7 +979,7 @@ impl<'b> CodeGenerator<'b> {
         UnsafeVarMarker::from_safe_vars(safe_vars)
     }
 
-    pub(crate) fn compile_fact(&mut self, term: &Term) -> Code {
+    pub(crate) fn compile_fact(&mut self, term: &Term) -> Result<Code, CompilationError> {
         self.update_var_count(post_order_iter(term));
 
         let mut vs = VariableFixtures::new();
@@ -993,6 +1001,10 @@ impl<'b> CodeGenerator<'b> {
                 false,
             );
 
+            if self.marker.max_reg_allocated() > MAX_ARITY {
+                return Err(CompilationError::from(ParserError::ExceededMaxArity));
+            }
+
             self.mark_unsafe_fact_vars(&mut compiled_fact);
 
             if !compiled_fact.is_empty() {
@@ -1001,7 +1013,7 @@ impl<'b> CodeGenerator<'b> {
         }
 
         code.push(instr!("proceed"));
-        code
+        Ok(code)
     }
 
     fn compile_query_line(
@@ -1111,7 +1123,7 @@ impl<'b> CodeGenerator<'b> {
             self.global_jmp_by_locs_offset = self.jmp_by_locs.len();
 
             let clause_code = match clause {
-                &PredicateClause::Fact(ref fact, ..) => self.compile_fact(fact),
+                &PredicateClause::Fact(ref fact, ..) => self.compile_fact(fact)?,
                 &PredicateClause::Rule(ref rule, ..) => self.compile_rule(rule)?,
             };
 

src/debray_allocator.rs

@@ -151,16 +151,13 @@ impl DebrayAllocator {
         };
     }
 
-    fn alloc_reg_to_var<'a, Target>(
+    fn alloc_reg_to_var<'a, Target: CompilationTarget<'a>>(
         &mut self,
         var: &String,
         lvl: Level,
         term_loc: GenContext,
         target: &mut Vec<Instruction>,
-    ) -> usize
+    ) -> usize {
-    where
-        Target: CompilationTarget<'a>,
-    {
         match term_loc {
             GenContext::Head => {
                 if let Level::Shallow = lvl {
@@ -404,4 +401,9 @@ impl Allocator for DebrayAllocator {
         self.arg_c = 1;
         self.temp_lb = arity + 1;
     }
+
+    #[inline(always)]
+    fn max_reg_allocated(&self) -> usize {
+        std::cmp::max(self.temp_lb, self.arg_c)
+    }
 }

src/machine/compile.rs

@@ -52,7 +52,7 @@ pub(super) fn compile_relation(
     match tl {
         &TopLevel::Query(_) => Err(CompilationError::ExpectedRel),
         &TopLevel::Predicate(ref clauses) => cg.compile_predicate(&clauses),
-        &TopLevel::Fact(ref fact, ..) => Ok(cg.compile_fact(fact)),
+        &TopLevel::Fact(ref fact, ..) => cg.compile_fact(fact),
         &TopLevel::Rule(ref rule, ..) => cg.compile_rule(rule),
     }
 }

src/machine/machine_errors.rs

@@ -585,16 +585,11 @@ impl MachineError {
 pub enum CompilationError {
     Arithmetic(ArithmeticError),
     ParserError(ParserError),
-    // BadPendingByte,
     CannotParseCyclicTerm,
-    // ExpandedTermsListNotAList,
     ExpectedRel,
-    // ExpectedTopLevelTerm,
     InadmissibleFact,
     InadmissibleQueryTerm,
     InconsistentEntry,
-    // InvalidDoubleQuotesDecl,
-    // InvalidHook,
     InvalidMetaPredicateDecl,
     InvalidModuleDecl,
     InvalidModuleExport,
@@ -631,18 +626,12 @@ impl CompilationError {
             &CompilationError::Arithmetic(..) => {
                 functor!(atom!("arithmetic_error"))
             }
-            // &CompilationError::BadPendingByte =>
-            //     functor!(atom_from_ss!("bad_pending_byte"), atom_tbl),
             &CompilationError::CannotParseCyclicTerm => {
                 functor!(atom!("cannot_parse_cyclic_term"))
             }
-            // &CompilationError::ExpandedTermsListNotAList =>
-            //     functor!(atom_tbl.build_with_static_str("expanded_terms_list_is_not_a_list")),
             &CompilationError::ExpectedRel => {
                 functor!(atom!("expected_relation"))
             }
-            // &CompilationError::ExpectedTopLevelTerm =>
-            //     functor!(atom_from_ss!("expected_atom_or_cons_or_clause"), atom_tbl),
             &CompilationError::InadmissibleFact => {
                 functor!(atom!("inadmissible_fact"))
             }
@@ -652,10 +641,6 @@ impl CompilationError {
             &CompilationError::InconsistentEntry => {
                 functor!(atom!("inconsistent_entry"))
             }
-            // &CompilationError::InvalidDoubleQuotesDecl =>
-            //     functor!(atom_from_ss!("invalid_double_quotes_declaration"), atom_tbl),
-            // &CompilationError::InvalidHook =>
-            //     functor!(atom_from_ss!("invalid_hook"), atom_tbl),
             &CompilationError::InvalidMetaPredicateDecl => {
                 functor!(atom!("invalid_meta_predicate_decl"))
             }

src/machine/preprocessor.rs

@@ -112,30 +112,6 @@ fn setup_predicate_indicator(term: &mut Term) -> Result<PredicateKey, Compilatio
     }
 }
 
-/*
-fn setup_scoped_predicate_indicator(term: &mut Term) -> Result<ScopedPredicateKey, CompilationError>
-{
-    match term {
-        Term::Clause(_, ref name, ref mut terms, Some(_))
-            if name.as_str() == ":" && terms.len() == 2 =>
-        {
-            let mut predicate_indicator = *terms.pop().unwrap();
-            let module_name = *terms.pop().unwrap();
-
-            let module_name = module_name
-                .to_constant()
-                .and_then(|c| c.to_atom())
-                .ok_or(CompilationError::InvalidModuleExport)?;
-
-            let key = setup_predicate_indicator(&mut predicate_indicator)?;
-
-            Ok((module_name, key))
-        }
-        _ => Err(CompilationError::InvalidModuleExport),
-    }
-}
-*/
-
 fn setup_module_export(
     mut term: Term,
     atom_tbl: &mut AtomTable,
@@ -279,6 +255,7 @@ fn setup_qualified_import(
  * -
  * ?
  */
+
 fn setup_meta_predicate<'a, LS: LoadState<'a>>(
     mut terms: Vec<Term>,
     loader: &mut Loader<'a, LS>,

src/parser/ast.rs

@@ -371,28 +371,29 @@ pub enum ArithmeticError {
 #[derive(Debug)]
 pub enum ParserError {
     BackQuotedString(usize, usize),
-    UnexpectedChar(char, usize, usize),
+    ExceededMaxArity,
-    UnexpectedEOF,
     IO(IOError),
     IncompleteReduction(usize, usize),
     InvalidSingleQuotedCharacter(char),
+    LexicalError(lexical::Error),
     MissingQuote(usize, usize),
     NonPrologChar(usize, usize),
     ParseBigInt(usize, usize),
-    LexicalError(lexical::Error),
+    UnexpectedChar(char, usize, usize),
+    UnexpectedEOF,
     Utf8Error(usize, usize),
 }
 
 impl ParserError {
     pub fn line_and_col_num(&self) -> Option<(usize, usize)> {
         match self {
-            &ParserError::BackQuotedString(line_num, col_num)
+            &ParserError::BackQuotedString(line_num, col_num) |
-            | &ParserError::UnexpectedChar(_, line_num, col_num)
+            &ParserError::IncompleteReduction(line_num, col_num) |
-            | &ParserError::IncompleteReduction(line_num, col_num)
+            &ParserError::MissingQuote(line_num, col_num) |
-            | &ParserError::MissingQuote(line_num, col_num)
+            &ParserError::NonPrologChar(line_num, col_num) |
-            | &ParserError::NonPrologChar(line_num, col_num)
+            &ParserError::ParseBigInt(line_num, col_num) |
-            | &ParserError::ParseBigInt(line_num, col_num)
+            &ParserError::UnexpectedChar(_, line_num, col_num) |
-            | &ParserError::Utf8Error(line_num, col_num) => Some((line_num, col_num)),
+            &ParserError::Utf8Error(line_num, col_num) => Some((line_num, col_num)),
             _ => None,
         }
     }
@@ -400,8 +401,7 @@ impl ParserError {
     pub fn as_atom(&self) -> Atom {
         match self {
             ParserError::BackQuotedString(..) => atom!("back_quoted_string"),
-            ParserError::UnexpectedChar(..) => atom!("unexpected_char"),
+            ParserError::ExceededMaxArity => atom!("exceeded_max_arity"),
-            ParserError::UnexpectedEOF => atom!("unexpected_end_of_file"),
             ParserError::IncompleteReduction(..) => atom!("incomplete_reduction"),
             ParserError::InvalidSingleQuotedCharacter(..) => atom!("invalid_single_quoted_character"),
             ParserError::IO(_) => atom!("input_output_error"),
@@ -409,6 +409,8 @@ impl ParserError {
             ParserError::MissingQuote(..) => atom!("missing_quote"),
             ParserError::NonPrologChar(..) => atom!("non_prolog_character"),
             ParserError::ParseBigInt(..) => atom!("cannot_parse_big_int"),
+            ParserError::UnexpectedChar(..) => atom!("unexpected_char"),
+            ParserError::UnexpectedEOF => atom!("unexpected_end_of_file"),
             ParserError::Utf8Error(..) => atom!("utf8_conversion_error"),
         }
     }

src/read.rs

@@ -52,7 +52,7 @@ impl MachineState {
         };
 
         inner.add_lines_read(num_lines_read);
-        Ok(write_term_to_heap(&term, &mut self.heap, &mut self.atom_tbl))
+        write_term_to_heap(&term, &mut self.heap, &mut self.atom_tbl)
     }
 }
 
@@ -245,7 +245,7 @@ pub(crate) fn write_term_to_heap(
     term: &Term,
     heap: &mut Heap,
     atom_tbl: &mut AtomTable,
-) -> TermWriteResult {
+) -> Result<TermWriteResult, ParserError> {
     let term_writer = TermWriter::new(heap, atom_tbl);
     term_writer.write_term_to_heap(term)
 }
@@ -311,7 +311,7 @@ impl<'a, 'b> TermWriter<'a, 'b> {
         }
     }
 
-    fn write_term_to_heap(mut self, term: &'a Term) -> TermWriteResult {
+    fn write_term_to_heap(mut self, term: &'a Term) -> Result<TermWriteResult, ParserError> {
         let heap_loc = self.heap.len();
 
         for term in breadth_first_iter(term, true) {
@@ -334,6 +334,10 @@ impl<'a, 'b> TermWriter<'a, 'b> {
                     self.push_stub_addr();
                 }
                 &TermRef::Clause(Level::Root, _, ref ct, subterms) => {
+                    if subterms.len() > MAX_ARITY {
+                        return Err(ParserError::ExceededMaxArity);
+                    }
+
                     self.heap.push(if subterms.len() == 0 {
                         heap_loc_as_cell!(heap_loc + 1)
                     } else {
@@ -413,9 +417,9 @@ impl<'a, 'b> TermWriter<'a, 'b> {
             self.modify_head_of_queue(&term, h);
         }
 
-        TermWriteResult {
+        Ok(TermWriteResult {
             heap_loc,
             var_dict: self.var_dict,
-        }
+        })
     }
 }