Repository: microsoft/rego-cpp (language: python)

Issue:
Policy not evaluated correctly when statements are not in order within comprehensions
This is similar to #101 
I came across this when trying to break the statement scheduling in Regorus, while using OPA and rego-cpp for comparison.

```rego
package test

x = [4, 5, 2, 3]

y = {a |
	b < 5
	a > 2
	x[a]
	x[b]
}

evaluates to `y = []` instead of `{3}`
```


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG b/CHANGELOG
index c7388483..74e8699b 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -1,5 +1,12 @@
 # Changelog
 
+## 2024-06-14 - Hardening
+Hardening for the 0.4.0 release.
+
+- Errors with individual input and data files are handled properly
+- Issues with binary inputs have been dealt with upstream in Trieste
+- Fixed an error with nested enumerations where results were incorrectly handled in some circumstances.
+
 ## 2024-06-12 - Update to rego v0.65.0
 Update to Rego v0.65.0
 
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 429e351c..3d073cd1 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -93,7 +93,7 @@ FetchContent_Declare(
 FetchContent_Declare(
   trieste
   GIT_REPOSITORY https://github.com/microsoft/trieste
-  GIT_TAG        c0234a5102fbc8ec84993dabcdbecedac82095d9
+  GIT_TAG        d7ccb706729e05b4a10eca91dac81f433470d9f7
 )
 
 FetchContent_MakeAvailable(cmake_utils)
@@ -199,8 +199,13 @@ endif()
 set(INSTALL_CONFIGDIR cmake)
 set(INSTALL_LIBDIR lib)
 set(INSTALL_INCLUDEDIR include)
+set(REGOCPP_INSTALL_TARGETS rego trieste json yaml snmalloc re2)
 
-install(TARGETS rego trieste json yaml snmalloc snmalloc-new-override re2
+if(TRIESTE_USE_SNMALLOC)
+  list(APPEND REGOCPP_INSTALL_TARGETS snmalloc-new-override)
+endif()
+
+install(TARGETS ${REGOCPP_INSTALL_TARGETS}
   EXPORT ${PROJECT_NAME}_Targets
   ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
   LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
@@ -240,7 +245,7 @@ endif()
 
 if (REGOCPP_BUILD_TESTS)
   install(TARGETS rego_test rego_test_c_api RUNTIME)
-  install(FILES tests/regocpp.yaml tests/bigint.yaml tests/bugs.yaml tests/regov1.yaml DESTINATION tests)
+  install(FILES tests/regocpp.yaml tests/bigint.yaml tests/bugs.yaml tests/regov1.yaml tests/bad.json DESTINATION tests)
   install(DIRECTORY tests/aci DESTINATION tests)
 endif()
 
diff --git a/src/interpreter.cc b/src/interpreter.cc
index 07ccf611..a1f0bb9d 100644
--- a/src/interpreter.cc
+++ b/src/interpreter.cc
@@ -271,8 +271,12 @@ namespace rego
         output_buf << "---" << std::endl;
         output_buf << "error: " << (error / ErrorMsg)->location().view()
                    << std::endl;
-        output_buf << "code: " << (error / ErrorCode)->location().view()
-                   << std::endl;
+        auto error_code = error->find_first(ErrorCode, error->begin());
+        if (error_code != error->end())
+        {
+          output_buf << "code: " << (*error_code)->location().view()
+                     << std::endl;
+        }
         output_buf << error_ast;
       }
     }
diff --git a/src/unifier.cc b/src/unifier.cc
index 6502304e..75247461 100644
--- a/src/unifier.cc
+++ b/src/unifier.cc
@@ -22,9 +22,8 @@ namespace rego
     m_version(version)
   {
     logging::Debug() << "---ASSEMBLING UNIFICATION---";
-    m_dependency_graph.push_back({"start", {}, 0});
+    m_dependency_graph.push_back({"start", {}});
     init_from_body(rulebody, m_statements, 0);
-    compute_dependency_scores();
     m_retries = detect_cycles();
     if (m_retries > 0)
     {
@@ -103,7 +102,7 @@ namespace rego
         std::size_t id = m_dependency_graph.size();
         statements.push_back({id, stmt});
         std::string name = "with" + std::to_string(id);
-        m_dependency_graph.push_back({name, {root}, 0});
+        m_dependency_graph.push_back({name, {root}});
         m_nested_statements[id] = {};
         init_from_body(stmt / UnifyBody, m_nested_statements[id], id);
 
@@ -128,7 +127,7 @@ namespace rego
         std::size_t id = m_dependency_graph.size();
         statements.push_back({id, stmt});
         std::string name = "not" + std::to_string(id);
-        m_dependency_graph.push_back({name, {root}, 0});
+        m_dependency_graph.push_back({name, {root}});
         m_nested_statements[id] = {};
         init_from_body(stmt->front(), m_nested_statements[id], root);
         for (auto& nested : m_nested_statements[id])
@@ -144,7 +143,7 @@ namespace rego
     auto name = (local / Var)->location();
     std::size_t id = m_dependency_graph.size();
     m_variables.insert({name, Variable(local, id)});
-    m_dependency_graph.push_back({std::string(name.view()), {}, 0});
+    m_dependency_graph.push_back({std::string(name.view()), {}});
     return id;
   }
 
@@ -172,64 +171,11 @@ namespace rego
     Resolver::expr_str(out, unifyexpr);
     std::string name = out.str();
     m_dependency_graph.push_back(
-      {name, std::set(dep_ids.begin(), dep_ids.end()), 0});
+      {name, std::set(dep_ids.begin(), dep_ids.end())});
     m_dependency_graph[var.id()].dependencies.insert(expr_id);
     return expr_id;
   }
 
-  std::size_t UnifierDef::compute_dependency_score(
-    std::size_t id, std::set<std::size_t>& visited)
-  {
-    if (contains(visited, id))
-    {
-      return m_dependency_graph[id].score;
-    }
-
-    visited.insert(id);
-    std::size_t score = 1;
-    for (auto dep_id : m_dependency_graph[id].dependencies)
-    {
-      score += compute_dependency_score(dep_id, visited);
-    }
-
-    m_dependency_graph[id].score = score;
-    return score;
-  }
-
-  void UnifierDef::compute_dependency_scores()
-  {
-    std::set<std::size_t> visited;
-    for (std::size_t id = 0; id < m_dependency_graph.size(); ++id)
-    {
-      compute_dependency_score(id, visited);
-    }
-
-    std::sort(
-      m_statements.begin(), m_statements.end(), [this](auto& lhs, auto& rhs) {
-        return dependency_score(lhs) < dependency_score(rhs);
-      });
-
-    for (auto& [_, nested_statements] : m_nested_statements)
-    {
-      std::sort(
-        nested_statements.begin(),
-        nested_statements.end(),
-        [this](auto& lhs, auto& rhs) {
-          return dependency_score(lhs) < dependency_score(rhs);
-        });
-    }
-  }
-
-  std::size_t UnifierDef::dependency_score(const Statement& stmt) const
-  {
-    return m_dependency_graph[stmt.id].score;
-  }
-
-  std::size_t UnifierDef::dependency_score(const Variable& var) const
-  {
-    return m_dependency_graph[var.id()].score;
-  }
-
   bool UnifierDef::is_local(const Node& var)
   {
     return is_variable(var->location());
diff --git a/src/unify.hh b/src/unify.hh
index 393fd637..a3284c72 100644
--- a/src/unify.hh
+++ b/src/unify.hh
@@ -610,7 +610,6 @@ namespace rego
     {
       std::string name;
       std::set<std::size_t> dependencies;
-      std::size_t score;
     };
 
     struct Statement
@@ -635,11 +634,6 @@ namespace rego
     void add_withpop(const Node& withpop);
     void reset();
 
-    void compute_dependency_scores();
-    std::size_t compute_dependency_score(
-      std::size_t index, std::set<size_t>& visited);
-    std::size_t dependency_score(const Variable& var) const;
-    std::size_t dependency_score(const Statement& stmt) const;
     std::size_t detect_cycles() const;
     bool has_cycle(std::size_t id) const;
 
@@ -705,7 +699,7 @@ namespace trieste::logging
     for (auto it = deps.begin(); it != deps.end(); ++it)
     {
       auto& dep = *it;
-      log << "[" << dep.name << "](" << dep.score << ") -> {";
+      log << "[" << dep.name << "] -> {";
 
       logging::Sep sep{", "};
       for (auto& idx : dep.dependencies)
diff --git a/src/unify/enumerate.cc b/src/unify/enumerate.cc
index 9817ffa0..476e7f5c 100644
--- a/src/unify/enumerate.cc
+++ b/src/unify/enumerate.cc
@@ -65,63 +65,6 @@ namespace
     auto it = unifybody->find(local) + 1;
     return find_enum(unifybody, it);
   }
-
-  void vars_from(Node node, std::set<Location>& vars)
-  {
-    if (node->type() == Var)
-    {
-      vars.insert(node->location());
-    }
-
-    for (Node child : *node)
-    {
-      vars_from(child, vars);
-    }
-  }
-
-  // Determines which statements following an implicit enum statement are needed
-  // to instantiate the item sequence for that enum and adds them to the outside
-  // list. Otherwise they should be captured by the enum and are placed in the
-  // inside list.
-  void capture_statements(
-    const NodeRange tail,
-    Node itemseq,
-    std::vector<Node>& outside,
-    std::vector<Node>& inside)
-  {
-    std::set<Location> vars;
-    vars_from(itemseq, vars);
-    for (Node stmt : tail)
-    {
-      if (stmt->type() == LiteralInit)
-      {
-        std::set<Location> init_vars;
-        vars_from(stmt / Lhs, init_vars);
-        vars_from(stmt / Rhs, init_vars);
-        std::set<Location> intersection;
-        std::set_intersection(
-          vars.begin(),
-          vars.end(),
-          init_vars.begin(),
-          init_vars.end(),
-          std::inserter(intersection, intersection.begin()));
-        if (intersection.empty())
-        {
-          inside.push_back(stmt);
-        }
-        else
-        {
-          // the item sequence depends on this init, so it must
-          // be outside the enum.
-          outside.push_back(stmt);
-        }
-      }
-      else
-      {
-        inside.push_back(stmt);
-      }
-    }
-  }
 }
 
 namespace rego
@@ -228,14 +171,10 @@ namespace rego
               return err(idx, "Invalid index for enumeration");
             }
 
-            std::vector<Node> outside;
-            std::vector<Node> inside;
-            capture_statements(_[Tail], _(ItemSeq), outside, inside);
-
             auto temp = _.fresh({"enum"});
             auto item = _.fresh({"item"});
             return Seq
-              << outside << (Local << (Var ^ item) << Undefined)
+              << (Local << (Var ^ item) << Undefined)
               << (LiteralEnum
                   << (Var ^ item) << _(ItemSeq)
                   << (UnifyBody
@@ -259,7 +198,7 @@ namespace rego
                                           << (Var ^ item)
                                           << (RefArgBrack
                                               << (Scalar << (Int ^ "1"))))))))
-                      << inside));
+                      << _[Tail]));
           },
 
         In(UnifyBody) *
@@ -301,14 +240,10 @@ namespace rego
               idx = Term << idx;
             }
 
-            std::vector<Node> outside;
-            std::vector<Node> inside;
-            capture_statements(_[Tail], _(ItemSeq), outside, inside);
-
             auto temp = _.fresh({"enum"});
             auto item = _.fresh({"item"});
             return Seq
-              << outside << (Local << (Var ^ item) << Undefined)
+              << (Local << (Var ^ item) << Undefined)
               << (LiteralEnum
                   << (Var ^ item) << _(ItemSeq)
                   << (UnifyBody
@@ -333,7 +268,7 @@ namespace rego
                                               << (RefArgBrack
                                                   << (Scalar
                                                       << (Int ^ "1")))))))))
-                      << inside));
+                      << _[Tail]));
           },
       }};
   }
diff --git a/src/unify/init.cc b/src/unify/init.cc
index 5b11095b..fb9ade61 100644
--- a/src/unify/init.cc
+++ b/src/unify/init.cc
@@ -2,24 +2,102 @@
 
 #include <algorithm>
 #include <cstddef>
-#include <deque>
+#include <list>
 
 namespace
 {
   using namespace rego;
   using namespace wf::ops;
 
-  struct InitSide
+  struct StmtSide
   {
     std::set<Location> vars;
     std::set<Location> inits;
   };
 
-  struct InitInfo
+  std::ostream& operator<<(std::ostream& os, const StmtSide& side)
+  {
+    os << "[{";
+    join(
+      os,
+      side.inits.begin(),
+      side.inits.end(),
+      ",",
+      [](auto& stream, auto& loc) {
+        stream << loc.view();
+        return true;
+      });
+    os << "} < {";
+    join(
+      os, side.vars.begin(), side.vars.end(), ",", [](auto& stream, auto& loc) {
+        stream << loc.view();
+        return true;
+      });
+    return os << "}]";
+  }
+
+  struct StmtInfo
   {
     std::size_t index;
-    InitSide lhs;
-    InitSide rhs;
+    StmtSide lhs;
+    StmtSide rhs;
+
+    bool remove(const Location& loc)
+    {
+      bool removed = lhs.vars.erase(loc) > 0;
+      removed |= lhs.inits.erase(loc) > 0;
+      removed |= rhs.vars.erase(loc) > 0;
+      removed |= rhs.inits.erase(loc) > 0;
+      return removed;
+    }
+
+    bool is_strict_init() const
+    {
+      return (!lhs.inits.empty() && rhs.vars.empty()) ||
+        (!rhs.inits.empty() && lhs.vars.empty());
+    }
+
+    bool is_init() const
+    {
+      return !lhs.inits.empty() || !rhs.inits.empty();
+    }
+  };
+
+  std::ostream& operator<<(std::ostream& os, const StmtInfo& info)
+  {
+    return os << info.index << ": " << info.lhs << " / " << info.rhs;
+  }
+
+  struct Locals
+  {
+    std::vector<std::set<Location>> locals;
+
+    void insert(const Location& loc)
+    {
+      locals.back().insert(loc);
+    }
+
+    void erase(const Location& loc)
+    {
+      locals.back().erase(loc);
+    }
+
+    bool contains(const Location& loc) const
+    {
+      return std::any_of(locals.rbegin(), locals.rend(), [loc](auto& set) {
+        return set.count(loc) > 0;
+      });
+    }
+
+    void push()
+    {
+      locals.push_back({});
+    }
+
+    void pop()
+    {
+      locals.pop_back();
+    }
   };
 
   Node to_init(
@@ -48,10 +126,9 @@ namespace
     return LiteralInit << lhs_vars << rhs_vars << (AssignInfix << lhs << rhs);
   }
 
-  void inits_from(
-    Node node, const std::set<Location>& locals, std::set<Location>& inits)
+  void inits_from(Node node, const Locals& locals, std::set<Location>& inits)
   {
-    if (node->type() == Var && contains(locals, node->location()))
+    if (node->type() == Var && locals.contains(node->location()))
     {
       inits.insert(node->location());
       return;
@@ -82,8 +159,7 @@ namespace
       node = node->front();
     }
 
-    std::set<Token> include_parents = {AssignArg, RefTerm, Object, Array};
-    if (!contains(include_parents, node->type()))
+    if (!node->in({AssignArg, RefTerm, Object, Array}))
     {
       return;
     }
@@ -94,29 +170,33 @@ namespace
     }
   }
 
-  void vars_from(
-    Node node, const std::set<Location>& locals, std::set<Location>& vars)
+  void vars_from(Node node, const Locals& locals, std::set<Location>& vars)
   {
-    if (node->type() == Var && contains(locals, node->location()))
+    if (node->type() == Var && locals.contains(node->location()))
     {
       vars.insert(node->location());
     }
 
+    if (node == RefArgDot)
+    {
+      return;
+    }
+
     for (Node child : *node)
     {
       vars_from(child, locals, vars);
     }
   }
 
-  InitSide side_from(Node node, const std::set<Location>& locals)
+  StmtSide side_from(Node node, Locals& locals)
   {
-    InitSide side;
+    StmtSide side;
     inits_from(node, locals, side.inits);
     vars_from(node, locals, side.vars);
     return side;
   }
 
-  bool any_compiler_inits(const InitSide& lhs)
+  bool any_compiler_inits(const StmtSide& lhs)
   {
     return std::any_of(lhs.inits.begin(), lhs.inits.end(), [](auto& loc) {
       std::string name = loc.str();
@@ -125,91 +205,57 @@ namespace
     });
   }
 
-  void remove_locals(
-    std::deque<InitInfo>& init_deque, const std::set<Location>& to_remove)
-  {
-    std::size_t count = init_deque.size();
-    for (std::size_t i = 0; i < count; ++i)
-    {
-      InitInfo& init_stmt = init_deque.front();
-      for (auto& loc : to_remove)
-      {
-        init_stmt.lhs.vars.erase(loc);
-        init_stmt.lhs.inits.erase(loc);
-        init_stmt.rhs.vars.erase(loc);
-        init_stmt.rhs.inits.erase(loc);
-      }
-      if (!init_stmt.lhs.inits.empty() || !init_stmt.rhs.inits.empty())
-      {
-        init_deque.push_back(init_stmt);
-      }
-
-      init_deque.pop_front();
-    }
-  }
-
-  std::vector<InitInfo> sort_init_stmts(
-    const std::set<Location>& locals, std::deque<InitInfo>& init_deque)
+  void find_init_stmts(Node unifybody, Locals& locals)
   {
-    std::set<Location> initialized;
-    std::vector<InitInfo> init_stmts;
-    while (!init_deque.empty() && initialized != locals)
-    {
-      // find all strict init statements
-      auto it =
-        std::find_if(init_deque.begin(), init_deque.end(), [](auto& init_stmt) {
-          return init_stmt.lhs.vars.empty() || init_stmt.rhs.vars.empty();
-        });
-
-      if (it == init_deque.end())
-      {
-        // we have a cycle, so we use the first statement
-        it = init_deque.begin();
-        init_stmts.push_back(*it);
-      }
-      else
-      {
-        init_stmts.push_back(*it);
-      }
-
-      std::set<Location> to_remove;
-      to_remove.insert(it->lhs.inits.begin(), it->lhs.inits.end());
-      to_remove.insert(it->rhs.inits.begin(), it->rhs.inits.end());
-      init_deque.erase(it);
-      remove_locals(init_deque, to_remove);
-      initialized.insert(to_remove.begin(), to_remove.end());
-    }
-
-    return init_stmts;
-  }
-
-  void find_init_stmts(Node unifybody, std::set<Location>& locals)
-  {
-    std::deque<InitInfo> potential_init_stmts;
+    logging::Trace() << std::endl
+                     << "Finding init statements and re-ordering body";
+    locals.push();
+    // Our ultimate goal is a re-ordering of statements such that all init
+    // statements come before the statements which depend on them
+    std::vector<StmtInfo> stmts;
+    std::map<Location, std::size_t> local_stmts;
     for (std::size_t i = 0; i < unifybody->size(); ++i)
     {
       Node stmt = unifybody->at(i);
       if (stmt->type() == Local)
       {
-        locals.insert((stmt / Var)->location());
+        // this is a local for this context
+        // as such it forms a starting point for the graph
+        Location loc = (stmt / Var)->location();
+        locals.insert(loc);
+        local_stmts[loc] = i;
+        stmts.push_back({i, {}, {}});
+        logging::Trace() << "Local: " << loc.view();
       }
       else if (stmt->type() == LiteralEnum)
       {
-        locals.erase((stmt / Item)->location());
+        // a literal enum initializes its item from the item sequence, which
+        // cannot be an init statement. The item variable is only
+        // used inside the enum, so we can safely ignore it.
+        Location item_loc = (stmt / Item)->location();
+        locals.erase(item_loc);
+        StmtSide rhs;
+        vars_from(stmt / ItemSeq, locals, rhs.vars);
+        stmts.push_back({i, {}, rhs});
       }
       else if (stmt->type() == Literal)
       {
         Node expr = stmt->front()->front();
         if (expr->type() != AssignInfix)
         {
+          // no possibility of assignment
+          // store this for later ordering
+          StmtSide rhs;
+          vars_from(stmt->front(), locals, rhs.vars);
+          stmts.push_back({i, {}, rhs});
           continue;
         }
 
         Node lhs = expr->front();
         Node rhs = expr->back();
 
-        InitSide lhs_side = side_from(lhs, locals);
-        InitSide rhs_side = side_from(rhs, locals);
+        StmtSide lhs_side = side_from(lhs, locals);
+        StmtSide rhs_side = side_from(rhs, locals);
 
         if (any_compiler_inits(lhs_side))
         {
@@ -218,74 +264,176 @@ namespace
           rhs_side.inits.clear();
         }
 
-        if (lhs_side.inits.empty() && rhs_side.inits.empty())
-        {
-          continue;
-        }
-
-        potential_init_stmts.push_back({i, lhs_side, rhs_side});
+        stmts.push_back({i, lhs_side, rhs_side});
+      }
+      else
+      {
+        StmtSide rhs;
+        vars_from(stmt->front(), locals, rhs.vars);
+        stmts.push_back({i, {}, rhs});
       }
     }
 
-    std::vector<InitInfo> init_stmts =
-      sort_init_stmts(locals, potential_init_stmts);
-    for (std::size_t i = 0; i < init_stmts.size(); ++i)
+    for (auto& stmt : stmts)
     {
-      InitInfo& init_stmt = init_stmts[i];
-      Node expr = unifybody->at(init_stmt.index)->front()->front();
+      logging::Trace() << stmt;
+    }
 
-      Node lhs = expr->front();
-      Node rhs = expr->back();
+    // build the graph
+    std::list<size_t> remaining(stmts.size());
+    std::iota(remaining.begin(), remaining.end(), 0);
+    std::vector<std::set<size_t>> edges(stmts.size());
+    while (!remaining.empty())
+    {
+      // find the first strict init statement
+      auto it =
+        std::find_if(remaining.begin(), remaining.end(), [stmts](auto& i) {
+          return stmts[i].is_strict_init();
+        });
 
-      for (auto& loc : init_stmt.lhs.inits)
+      if (it == remaining.end())
       {
-        locals.erase(loc);
+        // there is a loop, so break the tie using statement order
+        it = std::find_if(remaining.begin(), remaining.end(), [stmts](auto& i) {
+          return stmts[i].is_init();
+        });
+
+        if (it == remaining.end())
+        {
+          // no init statements left
+          break;
+        }
       }
 
-      for (auto& loc : init_stmt.rhs.inits)
+      size_t index = *it;
+      remaining.erase(it);
+
+      StmtInfo& stmt = stmts[index];
+      std::set<Location> init;
+      init.insert(stmt.lhs.inits.begin(), stmt.lhs.inits.end());
+      init.insert(stmt.rhs.inits.begin(), stmt.rhs.inits.end());
+      for (auto& loc : init)
       {
-        locals.erase(loc);
+        logging::Trace() << "removing init " << loc.view();
+        edges[stmt.index].insert(local_stmts[loc]);
+        for (auto& s : stmts)
+        {
+          if (s.index == index)
+          {
+            continue;
+          }
+
+          if (s.remove(loc))
+          {
+            edges[s.index].insert(index);
+          }
+        }
       }
+    }
 
-      unifybody->replace_at(
-        init_stmt.index,
-        to_init(lhs, init_stmt.lhs.inits, rhs, init_stmt.rhs.inits));
+    // topological sort
+    std::set<size_t> not_visited;
+    std::set<size_t> frontier;
+    for (size_t i = 0; i < stmts.size(); ++i)
+    {
+      not_visited.insert(i);
+      if (edges[i].empty())
+      {
+        frontier.insert(i);
+      }
     }
 
-    // where appropriate, recurse with the updated locals
-    for (Node stmt : *unifybody)
+    logging::Trace() << "Reordered statements:";
+    Nodes ordered_stmts;
+    while (!frontier.empty())
     {
-      if (stmt->type() == LiteralEnum)
+      size_t index = *frontier.begin();
+      frontier.erase(frontier.begin());
+      not_visited.erase(index);
+
+      StmtInfo& stmt = stmts[index];
+      logging::Trace() << stmt;
+      Node node = unifybody->at(index);
+      if (stmt.is_init())
       {
-        find_init_stmts(stmt / UnifyBody, locals);
+        Node expr = node->front()->front();
+
+        Node lhs = expr->front();
+        Node rhs = expr->back();
+        for (auto& loc : stmt.lhs.inits)
+        {
+          locals.erase(loc);
+        }
+
+        for (auto& loc : stmt.rhs.inits)
+        {
+          locals.erase(loc);
+        }
+
+        ordered_stmts.push_back(
+          to_init(lhs, stmt.lhs.inits, rhs, stmt.rhs.inits));
       }
-      else if (stmt->type() == LiteralWith)
+      else
       {
-        find_init_stmts(stmt / UnifyBody, locals);
+        ordered_stmts.push_back(node);
+      }
+
+      for (size_t i = 0; i < stmts.size(); ++i)
+      {
+        if (edges[i].erase(index) > 0 && edges[i].empty())
+        {
+          frontier.insert(i);
+        }
       }
-      else if (stmt->type() == LiteralNot)
+    }
+
+    if (!not_visited.empty())
+    {
+      logging::Warn() << "Unable to order body statements, possible error in "
+                         "unification logic";
+
+      while (!not_visited.empty())
+      {
+        size_t index = *not_visited.begin();
+        not_visited.erase(not_visited.begin());
+        Node node = unifybody->at(index);
+        ordered_stmts.push_back(node);
+      }
+    }
+
+    // where appropriate, recurse with the updated locals
+    for (Node stmt : ordered_stmts)
+    {
+      if (stmt->in({LiteralEnum, LiteralWith, LiteralNot}))
       {
         find_init_stmts(stmt / UnifyBody, locals);
       }
     }
+
+    unifybody->erase(unifybody->begin(), unifybody->end());
+    unifybody->insert(
+      unifybody->end(), ordered_stmts.begin(), ordered_stmts.end());
+    locals.pop();
   }
 
-  void register_init_stmts(const Node& rule)
+  int register_init_stmts(Node rule)
   {
     Node body = rule / Body;
     Node val = rule / Val;
 
     if (body->type() == UnifyBody)
     {
-      std::set<Location> locals;
+      Locals locals;
       find_init_stmts(body, locals);
     }
 
     if (val->type() == UnifyBody)
     {
-      std::set<Location> locals;
+      Locals locals;
       find_init_stmts(val, locals);
     }
+
+    return 0;
   }
 }
 
@@ -300,34 +448,18 @@ namespace rego
   {
     PassDef init = {"init", wf_pass_init, dir::once | dir::bottomup};
 
-    init.pre(RuleComp, [](Node rule) {
-      register_init_stmts(rule);
-      return 0;
-    });
-
-    init.pre(RuleFunc, [](Node rule) {
-      register_init_stmts(rule);
-      return 0;
-    });
-
-    init.pre(RuleSet, [](Node rule) {
-      register_init_stmts(rule);
-      return 0;
-    });
-
-    init.pre(RuleObj, [](Node rule) {
-      register_init_stmts(rule);
-      return 0;
-    });
-
+    init.pre(RuleComp, register_init_stmts);
+    init.pre(RuleFunc, register_init_stmts);
+    init.pre(RuleSet, register_init_stmts);
+    init.pre(RuleObj, register_init_stmts);
     init.pre(NestedBody, [](Node nested) {
-      std::set<Location> locals;
+      Locals locals;
       find_init_stmts(nested / Val, locals);
       return 0;
     });
 
     init.pre(ExprEvery, [](Node exprevery) {
-      std::set<Location> locals;
+      Locals locals;
       find_init_stmts(exprevery / UnifyBody, locals);
       return 0;
     });
diff --git a/tools/main.cc b/tools/main.cc
index 4c9bff71..84579385 100644
--- a/tools/main.cc
+++ b/tools/main.cc
@@ -64,7 +64,24 @@ int main(int argc, char** argv)
       return 1;
     }
 
-    interpreter.set_input_json_file(input_path);
+    rego::Node result;
+    try
+    {
+      result = interpreter.set_input_json_file(input_path);
+    }
+    catch (const std::exception& e)
+    {
+      trieste::logging::Error() << e.what() << std::endl;
+      return 1;
+    }
+
+    if (result != nullptr)
+    {
+      trieste::logging::Error()
+        << "Invalid input file: "
+        << std::filesystem::weakly_canonical(input_path) << std::endl;
+      return 1;
+    }
   }
 
   for (auto& path : data_paths)
@@ -76,13 +93,30 @@ int main(int argc, char** argv)
       return 1;
     }
 
-    if (path.extension() == ".json")
+    try
     {
-      interpreter.add_data_json_file(path);
+      rego::Node result;
+      if (path.extension() == ".json")
+      {
+        result = interpreter.add_data_json_file(path);
+      }
+      else
+      {
+        result = interpreter.add_module_file(path);
+      }
+
+      if (result != nullptr)
+      {
+        trieste::logging::Error()
+          << "Invalid data file:" << std::filesystem::weakly_canonical(path)
+          << std::endl;
+        return 1;
+      }
     }
-    else
+    catch (const std::exception& e)
     {
-      interpreter.add_module_file(path);
+      trieste::logging::Error() << e.what() << std::endl;
+      return 1;
     }
   }
 


Test patch:
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index b4f78787..3436bc33 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -13,12 +13,16 @@ if(REGOCPP_BUILD_TOOLS)
   add_test(NAME rego_fuzzer_unify COMMAND rego_fuzzer unify -f WORKING_DIRECTORY $<TARGET_FILE_DIR:rego_fuzzer>)
 endif()
 
+add_test(NAME rego_invalid_input COMMAND rego -i bad.json data WORKING_DIRECTORY $<TARGET_FILE_DIR:rego_interpreter>)
+add_test(NAME rego_invalid_large COMMAND rego -d rego data WORKING_DIRECTORY $<TARGET_FILE_DIR:rego_interpreter>)
 add_test(NAME rego_test_regocpp COMMAND rego_test regocpp.yaml -wf WORKING_DIRECTORY $<TARGET_FILE_DIR:rego_test>)
 add_test(NAME rego_test_bugs COMMAND rego_test bugs.yaml -wf WORKING_DIRECTORY $<TARGET_FILE_DIR:rego_test>)
 add_test(NAME rego_test_v1 COMMAND rego_test regov1.yaml -wf WORKING_DIRECTORY $<TARGET_FILE_DIR:rego_test>)
 add_test(NAME rego_test_aci COMMAND rego_test aci/aci.yaml WORKING_DIRECTORY $<TARGET_FILE_DIR:rego_test>)
 add_test(NAME rego_test_c_api COMMAND rego_test_c_api WORKING_DIRECTORY $<TARGET_FILE_DIR:rego_test>)
-set_tests_properties(rego_test_aci PROPERTIES TIMEOUT 300)
+set_property(TEST rego_invalid_input PROPERTY WILL_FAIL On)
+set_property(TEST rego_invalid_large PROPERTY WILL_FAIL On)
+set_property(TEST rego_test_aci PROPERTY TIMEOUT 300)
 
 add_custom_command(TARGET rego_test POST_BUILD
                    COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_CURRENT_SOURCE_DIR}/regocpp.yaml $<TARGET_FILE_DIR:rego_test>/regocpp.yaml)
@@ -26,6 +30,8 @@ add_custom_command(TARGET rego_test POST_BUILD
                    COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_CURRENT_SOURCE_DIR}/regov1.yaml $<TARGET_FILE_DIR:rego_test>/regov1.yaml)
 add_custom_command(TARGET rego_test POST_BUILD
                    COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_CURRENT_SOURCE_DIR}/bugs.yaml $<TARGET_FILE_DIR:rego_test>/bugs.yaml)
+add_custom_command(TARGET rego_test POST_BUILD
+                   COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_CURRENT_SOURCE_DIR}/bad.json $<TARGET_FILE_DIR:rego_interpreter>/bad.json)
 add_custom_command(TARGET rego_test POST_BUILD
                    COMMAND ${CMAKE_COMMAND} -E copy_directory ${CMAKE_CURRENT_SOURCE_DIR}/aci $<TARGET_FILE_DIR:rego_test>/aci)
 
diff --git a/tests/bad.json b/tests/bad.json
new file mode 100644
index 00000000..8b50e6e0
Binary files /dev/null and b/tests/bad.json differ
diff --git a/tests/bugs.yaml b/tests/bugs.yaml
index b8e500f4..e9a6e040 100644
--- a/tests/bugs.yaml
+++ b/tests/bugs.yaml
@@ -45,6 +45,32 @@ cases:
       - 1
       - 2
       - 3
+- note: regocpp/bug105
+  modules:
+  - |
+    package test
+
+    x = {5, 1, 2, 4, 3, 6}
+
+    p = {a |
+      b > a * 2
+      x[a]
+      x[b]
+    }
+
+    q = {a |
+      b := a / 3
+      x[a]
+      x[b]
+    }
+
+    y = [p, q]
+  query: data.test.y = x
+  want_result:
+    - x:
+      - [1, 2]
+      - [3, 6]
+
 - note: regocpp/bug110
   query: "[1, 2, 3][_] = x"
   want_result:
diff --git a/tests/test_case.cc b/tests/test_case.cc
index 9df6b995..cdca005c 100644
--- a/tests/test_case.cc
+++ b/tests/test_case.cc
@@ -406,6 +406,10 @@ namespace rego_test
     {
       if (wanted_bindings.find(key) == wanted_bindings.end())
       {
+        if (!pass)
+        {
+          os << "  ";
+        }
         os << "Actual binding: " << key << std::endl;
         pass = false;
       }