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

Issue:
OPA Test Driver
Write a driver that processes the OPA YAML-encoded test cases.


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/pr_gate.yml b/.github/workflows/pr_gate.yml
index 13b18128..f14a4935 100644
--- a/.github/workflows/pr_gate.yml
+++ b/.github/workflows/pr_gate.yml
@@ -8,6 +8,15 @@ on:
 jobs:
   cpp-format:
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        path: 
+          - check: src
+            exclude: ''
+          - check: test
+            exclude: ''
+          - check: tools
+            exclude: ''
     steps:
       - name: Checkout
         uses: actions/checkout@v3
@@ -16,8 +25,8 @@ jobs:
         uses: jidicula/clang-format-action@v4.8.0
         with:
           clang-format-version: '13'
-          check-path: src
-          exclude-regex: /examples/
+          check-path: ${{matrix.path['check']}}
+          exclude-regex: ${{matrix.path['exclude']}}
 
   linux:
     runs-on: ubuntu-latest
diff --git a/CHANGELOG b/CHANGELOG
index 31568e17..d91f3199 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -1,5 +1,18 @@
 # Changelog
 
+## [2023-08-12 - Version 0.2.4](https://github.com/microsoft/scenepic/releases/tag/v0.2.4)
+Adding a test driver that can execute OPA Rego-style YAML test files against the
+`regocpp` interpreter.
+
+**New Features**
+- The `rego_test` driver is a tool that can process OPA Rego-style test case YAML files.
+
+**Bug Fixes**
+- Fixed an issue with assignment of arrays and objects to local variables.
+
+**Improvements**
+- All existing tests have been converted to use the test driver
+
 ## [2023-08-08 - Version 0.2.3](https://github.com/microsoft/scenepic/releases/tag/v0.2.3)
 Adding the remainder of the missing language features, bringing the implementation fully
 compliant with the published v0.55.0 grammar.
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 4f00b8d3..b5b81064 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -49,6 +49,8 @@ else()
         src/*.h
         include/rego/*.h
         tools/*.cc
+        tests/*.cc
+        tests/*.h
   )
 
   add_custom_target(regocpp_format
@@ -69,6 +71,7 @@ option(REGOCPP_BUILD_TESTS "Specifies whether to build the tests" OFF)
 
 if(REGOCPP_BUILD_TESTS)
   enable_testing()
+  add_subdirectory(tests)
 endif()
 
 add_subdirectory(src)
@@ -118,6 +121,11 @@ if (REGOCPP_BUILD_TOOLS)
   install(DIRECTORY tools/examples DESTINATION .)
 endif()
 
+if (REGOCPP_BUILD_TESTS)
+  install(TARGETS rego_test yaml_trieste RUNTIME)
+  install(FILES tests/regocpp.yaml DESTINATION tests)
+endif()
+
 export(EXPORT ${PROJECT_NAME}_Targets
   FILE ${CMAKE_CURRENT_BINARY_DIR}/${PROJECT_NAME}Targets.cmake
   NAMESPACE ${PROJECT_NAME}::)
diff --git a/VERSION b/VERSION
index 373f8c6f..72f9fa82 100644
--- a/VERSION
+++ b/VERSION
@@ -1,1 +1,1 @@
-0.2.3
\ No newline at end of file
+0.2.4
\ No newline at end of file
diff --git a/include/rego/interpreter.h b/include/rego/interpreter.h
index cc910807..23315fc6 100644
--- a/include/rego/interpreter.h
+++ b/include/rego/interpreter.h
@@ -16,9 +16,12 @@ namespace rego
     void add_module(const std::string& name, const std::string& contents);
     void add_data_json_file(const std::filesystem::path& path);
     void add_data_json(const std::string& json);
+    void add_data(const Node& node);
     void add_input_json_file(const std::filesystem::path& path);
     void add_input_json(const std::string& json);
+    void add_input(const Node& node);
     std::string query(const std::string& query_expr) const;
+    Node raw_query(const std::string& query_expr) const;
     Interpreter& debug_path(const std::filesystem::path& prefix);
     const std::filesystem::path& debug_path() const;
     Interpreter& debug_enabled(bool enabled);
diff --git a/include/rego/lang.h b/include/rego/lang.h
index 8d14af7a..e86425d7 100644
--- a/include/rego/lang.h
+++ b/include/rego/lang.h
@@ -130,6 +130,7 @@ namespace rego
   inline const auto ArgVal = TokenDef("arg-val");
   inline const auto LiteralWith = TokenDef("literal-with");
   inline const auto LiteralEnum = TokenDef("literal-enum");
+  inline const auto LiteralInit = TokenDef("literal-init");
 
   // utility
   inline const auto Undefined = TokenDef("undefined");
diff --git a/include/rego/passes.h b/include/rego/passes.h
index a78f38a7..990e3ef3 100644
--- a/include/rego/passes.h
+++ b/include/rego/passes.h
@@ -53,6 +53,7 @@ namespace rego
   PassDef skip_refs();
   PassDef simple_refs();
   PassDef implicit_enums();
+  PassDef init();
   PassDef rulebody();
   PassDef lift_to_rule();
   PassDef functions();
diff --git a/include/rego/rego.h b/include/rego/rego.h
new file mode 100644
index 00000000..c2ab08c6
--- /dev/null
+++ b/include/rego/rego.h
@@ -0,0 +1,4 @@
+#pragma once
+
+#include "interpreter.h"
+#include "log.h"
\ No newline at end of file
diff --git a/include/rego/unifier.h b/include/rego/unifier.h
index d6eebe0c..e5740cd4 100644
--- a/include/rego/unifier.h
+++ b/include/rego/unifier.h
@@ -15,38 +15,38 @@ namespace rego
 {
   using namespace trieste;
 
+  class UnifierDef;
+  using Unifier = std::shared_ptr<UnifierDef>;
   using CallStack = std::shared_ptr<std::vector<Location>>;
   using ValuesLookup = std::map<std::string, Values>;
   using WithStack = std::shared_ptr<std::vector<ValuesLookup>>;
+  using UnifierCache = std::shared_ptr<NodeMap<Unifier>>;
 
-  class Unifier
+  class UnifierDef
   {
   public:
-    Unifier(
+    UnifierDef(
       const Location& rule,
       const Node& rulebody,
       CallStack call_stack,
       WithStack with_stack,
       const BuiltIns& builtins,
-      std::size_t max_passes = 100);
-    ~Unifier();
+      UnifierCache cache);
     Node unify();
     Nodes expressions() const;
     Nodes bindings() const;
     std::string str() const;
     std::string dependency_str() const;
-    friend std::ostream& operator<<(std::ostream& os, const Unifier& unifier);
-
-  private:
-    Unifier(
+    static Unifier create(
       const Location& rule,
       const Node& rulebody,
       CallStack call_stack,
       WithStack with_stack,
       const BuiltIns& builtins,
-      std::size_t max_passes,
-      Unifier* outer);
-    std::shared_ptr<Unifier> nested(const Location& rule, const Node& rulebody);
+      UnifierCache cache);
+    friend std::ostream& operator<<(std::ostream& os, const Unifier& unifier);
+
+  private:
     Unifier rule_unifier(const Location& rule, const Node& rulebody);
     void init_from_body(const Node& rulebody, std::vector<Node>& statements);
     void add_variable(const Node& local);
@@ -98,9 +98,8 @@ namespace rego
     CallStack m_call_stack;
     WithStack m_with_stack;
     BuiltIns m_builtins;
-    std::size_t m_max_passes;
     std::size_t m_retries;
     Token m_parent_type;
-    Unifier* m_outer;
+    UnifierCache m_cache;
   };
 }
\ No newline at end of file
diff --git a/include/rego/wf.h b/include/rego/wf.h
index 373fcd28..89df122d 100644
--- a/include/rego/wf.h
+++ b/include/rego/wf.h
@@ -441,8 +441,16 @@ namespace rego
   inline const auto wf_pass_implicit_enums = wf_pass_simple_refs;
 
   // clang-format off
-  inline const auto wf_pass_rulebody =
+  inline const auto wf_pass_init =
     wf_pass_implicit_enums
+    | (UnifyBody <<= (Local | Literal | LiteralWith | LiteralEnum | LiteralInit)++[1])
+    | (LiteralInit <<= AssignInfix)
+    ;
+  // clang-format on
+
+  // clang-format off
+  inline const auto wf_pass_rulebody =
+    wf_pass_init
     | (Module <<= (Import | RuleComp | DefaultRule | RuleFunc | RuleSet | RuleObj | Submodule)++)
     | (UnifyExpr <<= Var * (Val >>= NotExpr | Expr))
     | (Expr <<= NumTerm | RefTerm | Term | UnaryExpr | ArithInfix | BinInfix | BoolInfix | ExprCall | ExprEvery | Enumerate)
@@ -488,16 +496,16 @@ namespace rego
   // clang-format off
   inline const auto wf_pass_unify =
     wf_pass_functions
-    | (Query <<= (Term | Binding)++[1])
+    | (Query <<= (Term | Binding | Undefined)++[1])
     | (Binding <<= Var * Term)[Var]
-    | (Term <<= Scalar | Array | Object | Set | Undefined)
+    | (Term <<= Scalar | Array | Object | Set)
     ;
   // clang-format on
 
   // clang-format off
   inline const auto wf_pass_query =
     wf_pass_unify
-    | (Top <<= (Binding | Term)++[1])
+    | (Top <<= (Binding | Term | Undefined)++[1])
     ;
   // clang-format on
 }
\ No newline at end of file
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index eda4250a..354699ca 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -40,6 +40,7 @@ passes/comparison.cc
 passes/assign.cc
 passes/skip_refs.cc
 passes/simple_refs.cc
+passes/init.cc
 passes/rulebody.cc
 passes/lift_to_rule.cc
 passes/functions.cc
diff --git a/src/interpreter.cc b/src/interpreter.cc
index 7af8ad85..1b2924e7 100644
--- a/src/interpreter.cc
+++ b/src/interpreter.cc
@@ -64,6 +64,11 @@ namespace rego
     m_data_seq->push_back(data);
   }
 
+  void Interpreter::add_data(const Node& node)
+  {
+    m_data_seq->push_back(node);
+  }
+
   void Interpreter::add_input_json_file(const std::filesystem::path& path)
   {
     if (m_input->size() > 0)
@@ -92,6 +97,16 @@ namespace rego
     m_input->push_back(input);
   }
 
+  void Interpreter::add_input(const Node& node)
+  {
+    if (m_input->size() > 0)
+    {
+      throw std::runtime_error("Input already set");
+    }
+
+    m_input->push_back(node);
+  }
+
   bool Interpreter::has_error(const Node& node) const
   {
     if (node->type() == Error)
@@ -110,7 +125,7 @@ namespace rego
     return false;
   }
 
-  std::string Interpreter::query(const std::string& query_expr) const
+  Node Interpreter::raw_query(const std::string& query_expr) const
   {
     auto ast = NodeDef::create(Top);
     auto rego = NodeDef::create(rego::Rego);
@@ -182,6 +197,12 @@ namespace rego
       }
     }
 
+    return ast;
+  }
+
+  std::string Interpreter::query(const std::string& query_expr) const
+  {
+    Node ast = raw_query(query_expr);
     std::ostringstream result_buf;
     for (auto result : *ast)
     {
diff --git a/src/lang.cc b/src/lang.cc
index f2c9c38f..2b5119e0 100644
--- a/src/lang.cc
+++ b/src/lang.cc
@@ -61,6 +61,7 @@ namespace rego
       {"skip_refs", skip_refs(), &wf_pass_skip_refs},
       {"simple_refs", simple_refs(), &wf_pass_simple_refs},
       {"implicit_enums", implicit_enums(), &wf_pass_implicit_enums},
+      {"init", init(), &wf_pass_init},
       {"rulebody", rulebody(), &wf_pass_rulebody},
       {"lift_to_rule", lift_to_rule(), &wf_pass_lift_to_rule},
       {"functions", functions(), &wf_pass_functions},
@@ -71,7 +72,6 @@ namespace rego
 
   Driver& driver()
   {
-    auto passdefs = passes();
     static Driver d(
       "rego",
       nullptr,
@@ -109,6 +109,7 @@ namespace rego
         {"skip_refs", skip_refs(), wf_pass_skip_refs},
         {"simple_refs", simple_refs(), wf_pass_simple_refs},
         {"implicit_enums", implicit_enums(), wf_pass_implicit_enums},
+        {"init", init(), wf_pass_init},
         {"rulebody", rulebody(), wf_pass_rulebody},
         {"lift_to_rule", lift_to_rule(), wf_pass_lift_to_rule},
         {"functions", functions(), wf_pass_functions},
diff --git a/src/parse.cc b/src/parse.cc
index 3c376d37..b4800351 100644
--- a/src/parse.cc
+++ b/src/parse.cc
@@ -6,6 +6,33 @@ namespace rego
     "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
   const std::string numeric = "0123456789";
   const std::string alphanumeric = alpha + numeric;
+  const std::size_t max_string_length = 10;
+
+  template <typename T>
+  std::string rand_string(T& rnd)
+  {
+    // TODO add valid non-alphanum characters
+    std::ostringstream buf;
+    std::size_t length = rnd() % max_string_length;
+    for (std::size_t i = 0; i < length; ++i)
+    {
+      buf << alphanumeric[rnd() % alphanumeric.size()];
+    }
+    return buf.str();
+  }
+
+  template <typename T>
+  std::string rand_raw_string(T& rnd)
+  {
+    // TODO add valid non-alphanum characters
+    std::ostringstream buf;
+    std::size_t length = rnd() % max_string_length;
+    for (std::size_t i = 0; i < length; ++i)
+    {
+      buf << alphanumeric[rnd() % alphanumeric.size()];
+    }
+    return buf.str();
+  }
 
   Parse parser()
   {
@@ -217,6 +244,8 @@ namespace rego
       JSONTrue >> [](auto&) { return "true"; },
       JSONFalse >> [](auto&) { return "false"; },
       JSONNull >> [](auto&) { return "null"; },
+      JSONString >> [](auto& rnd) { return rand_string(rnd); },
+      RawString >> [](auto& rnd) { return rand_raw_string(rnd); },
     });
 
     return p;
diff --git a/src/passes/init.cc b/src/passes/init.cc
new file mode 100644
index 00000000..995df001
--- /dev/null
+++ b/src/passes/init.cc
@@ -0,0 +1,130 @@
+#include "lang.h"
+#include "log.h"
+#include "passes.h"
+
+namespace
+{
+  using namespace rego;
+  using namespace wf::ops;
+
+  using InitStatments = std::shared_ptr<NodeMap<Location>>;
+
+  void vars_from(Node node, std::set<Location>& vars)
+  {
+    std::set<Token> exclude_parents = {RefArgDot, VarSeq, UnifyBody};
+    if (exclude_parents.contains(node->type()))
+    {
+      return;
+    }
+
+    if (node->type() == Var)
+    {
+      vars.insert(node->location());
+    }
+    else
+    {
+      for (Node child : *node)
+      {
+        vars_from(child, vars);
+      }
+    }
+  }
+
+  void find_init_stmts(Node unifybody, InitStatments init_stmts)
+  {
+    std::set<Location> locals;
+    std::set<Location> initialized;
+    for (Node stmt : *unifybody)
+    {
+      if (stmt->type() == Local)
+      {
+        locals.insert((stmt / Var)->location());
+      }
+      else if (stmt->type() == Literal)
+      {
+        std::set<Location> vars;
+        vars_from(stmt, vars);
+        std::set<Location> found;
+        std::set_intersection(
+          vars.begin(),
+          vars.end(),
+          locals.begin(),
+          locals.end(),
+          std::inserter(found, found.begin()));
+        if (found.size() == 1)
+        {
+          Location loc = *found.begin();
+          initialized.insert(loc);
+          locals.erase(loc);
+          init_stmts->insert({stmt, loc});
+        }
+      }
+    }
+  }
+
+  bool contains_loc(const Node& node, const Location& loc)
+  {
+    std::set<Token> exclude_parents = {RefArgDot, VarSeq, UnifyBody};
+    if (exclude_parents.contains(node->type()))
+    {
+      return false;
+    }
+
+    if (node->type() == Var)
+    {
+      return node->location() == loc;
+    }
+    else
+    {
+      for (Node child : *node)
+      {
+        if (contains_loc(child, loc))
+        {
+          return true;
+        }
+      }
+      return false;
+    }
+  }
+}
+
+namespace rego
+{
+  using namespace wf::ops;
+
+  // Convert all Literal nodes into UnifyExpr nodes of the form <var> = <expr> |
+  // <not-expr>.
+  PassDef init()
+  {
+    InitStatments init_stmts = std::make_shared<NodeMap<Location>>();
+
+    PassDef init = {
+      dir::once | dir::topdown,
+      {
+        In(UnifyBody) *
+            (T(Literal)[Literal](
+               [init_stmts](auto& n) { return init_stmts->contains(*n.first); })
+             << (T(Expr)
+                 << (T(AssignInfix)
+                     << (T(AssignArg)[Lhs] * T(AssignArg)[Rhs])))) >>
+          [init_stmts](Match& _) {
+            Location init_var = init_stmts->at(_(Literal));
+            if (contains_loc(_(Lhs), init_var))
+            {
+              return LiteralInit << (AssignInfix << _(Lhs) << _(Rhs));
+            }
+            else
+            {
+              return LiteralInit << (AssignInfix << _(Rhs) << _(Lhs));
+            }
+          },
+      }};
+
+    init.pre(UnifyBody, [init_stmts](Node unifybody) {
+      find_init_stmts(unifybody, init_stmts);
+      return 0;
+    });
+
+    return init;
+  }
+}
\ No newline at end of file
diff --git a/src/passes/lists.cc b/src/passes/lists.cc
index f28e8497..5ca2c65f 100644
--- a/src/passes/lists.cc
+++ b/src/passes/lists.cc
@@ -21,6 +21,11 @@ namespace
 
   bool contains(const Node& node, const std::set<Token>& tokens)
   {
+    if (node->type() == Error)
+    {
+      return false;
+    }
+
     if (tokens.contains(node->type()))
     {
       return true;
diff --git a/src/passes/rulebody.cc b/src/passes/rulebody.cc
index 74635580..0a86ad0f 100644
--- a/src/passes/rulebody.cc
+++ b/src/passes/rulebody.cc
@@ -50,11 +50,23 @@ namespace rego
                                        << _(UnifyBody));
         },
 
+      // <var> := <any>
+      In(UnifyBody) *
+          (T(LiteralInit)
+           << (T(AssignInfix)
+               << ((T(AssignArg) << (T(RefTerm) << T(Var)[Lhs])) *
+                   T(AssignArg)[Rhs]))) >>
+        [](Match& _) {
+          LOG("<var> := <any>");
+          return UnifyExpr << _(Lhs) << (Expr << _(Rhs)->front());
+        },
+
       In(UnifyBody) *
           (T(Literal) << ((T(Expr) << T(AssignInfix)[AssignInfix]))) >>
         [](Match& _) { return _(AssignInfix); },
 
-      In(With) * T(Expr)[Expr] >>
+      In(With) *
+          T(Expr)[Expr]([](auto& n) { return is_in(*n.first, {UnifyBody}); }) >>
         [](Match& _) {
           LOG("with");
           Location temp = _.fresh({"with"});
@@ -136,7 +148,7 @@ namespace rego
         },
 
       // <array> = <array>
-      In(UnifyBody) *
+      (In(UnifyBody) / In(LiteralInit)) *
           (T(AssignInfix)
            << ((T(AssignArg) << (T(Term) << T(Array)[Lhs])) *
                (T(AssignArg) << (T(Term) << T(Array)[Rhs])))) >>
@@ -160,7 +172,7 @@ namespace rego
         },
 
       // <object> = <object>
-      In(UnifyBody) *
+      (In(UnifyBody) / In(LiteralInit)) *
           (T(AssignInfix)
            << ((T(AssignArg) << (T(Term) << T(Object)[Lhs])) *
                (T(AssignArg) << (T(Term) << T(Object)[Rhs])))) >>
@@ -201,7 +213,7 @@ namespace rego
         },
 
       // <array> = <var>
-      In(UnifyBody) *
+      (In(UnifyBody) / In(LiteralInit)) *
           (T(AssignInfix)
            << ((T(AssignArg) << (T(Term) << T(Array)[Lhs])) *
                (T(AssignArg) << (T(RefTerm) << T(Var)[Rhs])))) >>
@@ -221,7 +233,7 @@ namespace rego
         },
 
       // <array> = <ref>
-      In(UnifyBody) *
+      (In(UnifyBody) / In(LiteralInit)) *
           (T(AssignInfix)
            << ((T(AssignArg) << (T(Term) << T(Array)[Lhs])) *
                (T(AssignArg) << T(RefTerm)[RefTerm]))) >>
@@ -256,7 +268,7 @@ namespace rego
         },
 
       // <object> = <var>
-      In(UnifyBody) *
+      (In(UnifyBody) / In(LiteralInit)) *
           (T(AssignInfix)
            << ((T(AssignArg) << (T(Term) << T(Object)[Lhs])) *
                (T(AssignArg) << (T(RefTerm) << T(Var)[Rhs])))) >>
@@ -298,7 +310,7 @@ namespace rego
         },
 
       // <object> = <ref>
-      In(UnifyBody) *
+      (In(UnifyBody) / In(LiteralInit)) *
           (T(AssignInfix)
            << ((T(AssignArg) << (T(Term) << T(Object)[Lhs])) *
                (T(AssignArg) << T(RefTerm)[RefTerm]))) >>
@@ -336,6 +348,9 @@ namespace rego
           return seq;
         },
 
+      In(UnifyBody) * (T(LiteralEnum) << T(UnifyExpr)[UnifyExpr]) >>
+        [](Match& _) { return _(UnifyExpr); },
+
       // <numbery> = <numbery>
       In(UnifyBody) *
           (T(AssignInfix)
@@ -413,7 +428,9 @@ namespace rego
         },
 
       // <binarg>.<setcompr>
-      In(BinArg) * T(SetCompr)[SetCompr] >>
+      In(BinArg) * T(SetCompr)[SetCompr]([](auto& n) {
+        return is_in(*n.first, {UnifyBody});
+      }) >>
         [](Match& _) {
           LOG("<binarg>.<setcompr>");
           Location set = _.fresh({"compr"});
@@ -434,6 +451,22 @@ namespace rego
 
       (In(UnifyBody) / In(Expr)) * T(AssignInfix)[AssignInfix] >>
         [](Match& _) { return err(_(AssignInfix), "Invalid assignment"); },
+
+      In(UnifyBody) * (T(LiteralInit) << T(AssignInfix)[AssignInfix]) >>
+        [](Match& _) { return err(_(AssignInfix), "Invalid assignment"); },
+
+      (In(Term) / In(BinArg)) * T(SetCompr)[SetCompr] >>
+        [](Match& _) { return err(_(SetCompr), "Invalid set comprehension"); },
+
+      In(Term) * T(ArrayCompr)[ArrayCompr] >>
+        [](Match& _) {
+          return err(_(ArrayCompr), "Invalid array comprehension");
+        },
+
+      In(Term) * T(ObjectCompr)[ObjectCompr] >>
+        [](Match& _) {
+          return err(_(ObjectCompr), "Invalid object comprehension");
+        },
     };
   }
 }
\ No newline at end of file
diff --git a/src/passes/symbols.cc b/src/passes/symbols.cc
index 123d30df..578f876f 100644
--- a/src/passes/symbols.cc
+++ b/src/passes/symbols.cc
@@ -12,7 +12,7 @@ namespace
 
   void find_assigned_vars(const Node& node, Nodes& vars)
   {
-    if (node->type() == LiteralWith)
+    if (node->type() == LiteralWith || node->type() == SomeDecl)
     {
       return;
     }
@@ -271,6 +271,7 @@ namespace rego
         [](Match& _) {
           Location item = _.fresh({"item"});
           Location every = _.fresh({"every"});
+
           return Seq
             << (Lift << UnifyBody << (Local << (Var ^ item) << Undefined))
             << (Lift << UnifyBody << (LiteralEnum << (Var ^ item) << _(Expr)))
diff --git a/src/resolver.cc b/src/resolver.cc
index 8ac3e5d7..ee9a9f23 100644
--- a/src/resolver.cc
+++ b/src/resolver.cc
@@ -1111,13 +1111,14 @@ namespace rego
 
     Node rulebody = defs[0] / Val;
     {
-      Unifier unifier(
+      auto unifier = UnifierDef::create(
         {"query"},
         rulebody,
         std::make_shared<std::vector<Location>>(),
         std::make_shared<std::vector<ValuesLookup>>(),
-        BuiltIns().register_standard_builtins());
-      unifier.unify();
+        BuiltIns().register_standard_builtins(),
+        std::make_shared<NodeMap<Unifier>>());
+      unifier->unify();
     }
 
     Node result = NodeDef::create(Query);
@@ -1150,6 +1151,11 @@ namespace rego
         }
       }
 
+      if (term->type() == Undefined)
+      {
+        continue;
+      }
+
       if (term->type() != Term)
       {
         term = Term << term;
@@ -1169,6 +1175,11 @@ namespace rego
       }
     }
 
+    if (result->size() == 0)
+    {
+      result->push_back(Undefined);
+    }
+
     return result;
   }
 }
\ No newline at end of file
diff --git a/src/unifier.cc b/src/unifier.cc
index dcb2c2c4..6dd61043 100644
--- a/src/unifier.cc
+++ b/src/unifier.cc
@@ -42,7 +42,7 @@ namespace rego
 
   std::ostream& operator<<(std::ostream& os, const Unifier& unifier)
   {
-    for (const auto& [first, second] : unifier.m_variables)
+    for (const auto& [first, second] : unifier->m_variables)
     {
       os << first.view() << "(" << second.dependency_score() << ") -> {";
       std::string sep = "";
@@ -54,7 +54,7 @@ namespace rego
       os << "}" << std::endl;
     }
     os << std::endl;
-    for (const auto& unifyexpr : unifier.m_statements)
+    for (const auto& unifyexpr : unifier->m_statements)
     {
       if (unifyexpr->type() == UnifyExprWith)
       {
@@ -68,53 +68,30 @@ namespace rego
     return os;
   }
 
-  Unifier::Unifier(
+  UnifierDef::UnifierDef(
     const Location& rule,
     const Node& rulebody,
     CallStack call_stack,
     WithStack with_stack,
     const BuiltIns& builtins,
-    std::size_t max_passes) :
-    Unifier(
-      rule, rulebody, call_stack, with_stack, builtins, max_passes, nullptr)
-  {}
-
-  Unifier::Unifier(
-    const Location& rule,
-    const Node& rulebody,
-    CallStack call_stack,
-    WithStack with_stack,
-    const BuiltIns& builtins,
-    std::size_t max_passes,
-    Unifier* outer) :
+    UnifierCache cache) :
     m_rule(rule),
     m_call_stack(call_stack),
     m_with_stack(with_stack),
     m_builtins(builtins),
-    m_max_passes(max_passes),
     m_parent_type(rulebody->parent()->type()),
-    m_outer(outer)
+    m_cache(cache)
   {
-    if (this->push_rule(rule))
-    {
-      throw std::runtime_error("Recursion detected in rule body");
-    }
     init_from_body(rulebody, m_statements);
     compute_dependency_scores();
     m_retries = Variable::detect_cycles(m_variables);
     if (m_retries > 0)
     {
       LOG("Detected ", m_retries, " cycles in dependency graph");
-      m_max_passes = std::max(m_max_passes, m_retries + 1);
     }
   }
 
-  Unifier::~Unifier()
-  {
-    this->pop_rule(m_rule);
-  }
-
-  void Unifier::init_from_body(
+  void UnifierDef::init_from_body(
     const Node& rulebody, std::vector<Node>& statements)
   {
     std::for_each(rulebody->begin(), rulebody->end(), [&](const auto& stmt) {
@@ -136,13 +113,13 @@ namespace rego
     });
   }
 
-  void Unifier::add_variable(const Node& local)
+  void UnifierDef::add_variable(const Node& local)
   {
     auto name = (wfi / local / Var)->location();
     m_variables.insert({name, Variable(local)});
   }
 
-  void Unifier::add_unifyexpr(const Node& unifyexpr)
+  void UnifierDef::add_unifyexpr(const Node& unifyexpr)
   {
     Node lhs = wfi / unifyexpr / Var;
     Node rhs = wfi / unifyexpr / Val;
@@ -160,7 +137,7 @@ namespace rego
     var.insert_dependencies(deps.begin(), deps.end());
   }
 
-  void Unifier::compute_dependency_scores()
+  void UnifierDef::compute_dependency_scores()
   {
     std::set<Location> visited;
     for (auto& [_, variable] : m_variables)
@@ -184,7 +161,7 @@ namespace rego
     }
   }
 
-  std::size_t Unifier::compute_dependency_score(
+  std::size_t UnifierDef::compute_dependency_score(
     Variable& variable, std::set<Location>& visited)
   {
     if (visited.contains(variable.name()))
@@ -206,12 +183,12 @@ namespace rego
     return score;
   }
 
-  bool Unifier::is_local(const Node& var)
+  bool UnifierDef::is_local(const Node& var)
   {
     return is_variable(var->location());
   }
 
-  std::size_t Unifier::scan_vars(
+  std::size_t UnifierDef::scan_vars(
     const Node& expr, std::vector<Location>& locals)
   {
     std::size_t num_vars = 0;
@@ -239,7 +216,8 @@ namespace rego
     return num_vars;
   }
 
-  void Unifier::execute_statements(Nodes::iterator begin, Nodes::iterator end)
+  void UnifierDef::execute_statements(
+    Nodes::iterator begin, Nodes::iterator end)
   {
     for (auto it = begin; it != end; ++it)
     {
@@ -271,7 +249,7 @@ namespace rego
     }
   }
 
-  void Unifier::reset()
+  void UnifierDef::reset()
   {
     for (auto& [_, var] : m_variables)
     {
@@ -279,13 +257,13 @@ namespace rego
     }
   }
 
-  void Unifier::pass()
+  void UnifierDef::pass()
   {
     LOG_MAP_VALUES(m_variables);
     execute_statements(m_statements.begin(), m_statements.end());
   }
 
-  void Unifier::mark_invalid_values()
+  void UnifierDef::mark_invalid_values()
   {
     for (auto& [_, var] : m_variables)
     {
@@ -299,14 +277,14 @@ namespace rego
     }
   }
 
-  void Unifier::remove_invalid_values()
+  void UnifierDef::remove_invalid_values()
   {
     std::for_each(m_variables.begin(), m_variables.end(), [](auto& pair) {
       pair.second.remove_invalid_values();
     });
   }
 
-  Node Unifier::bind_variables()
+  Node UnifierDef::bind_variables()
   {
     return std::accumulate(
       m_variables.begin(),
@@ -331,9 +309,13 @@ namespace rego
       });
   }
 
-  Node Unifier::unify()
+  Node UnifierDef::unify()
   {
-    reset();
+    if (this->push_rule(this->m_rule))
+    {
+      throw std::runtime_error("Recursion detected in rule body");
+    }
+
     LOG_HEADER("Unification", "=====");
 
     LOG("exprs: ");
@@ -354,10 +336,11 @@ namespace rego
     LOG_HEADER("Complete", "=====");
 
     Node result = bind_variables();
+    this->pop_rule(this->m_rule);
     return result;
   }
 
-  Args Unifier::create_args(const Node& args)
+  Args UnifierDef::create_args(const Node& args)
   {
     Args function_args;
     for (auto arg : *args)
@@ -382,7 +365,7 @@ namespace rego
     return function_args;
   }
 
-  Values Unifier::evaluate_function(
+  Values UnifierDef::evaluate_function(
     const Location& var, const std::string& func_name, const Args& args)
   {
     LOG("> calling ", func_name, " with ", args);
@@ -410,7 +393,7 @@ namespace rego
     return values;
   }
 
-  Values Unifier::evaluate(const Location& var, const Node& value)
+  Values UnifierDef::evaluate(const Location& var, const Node& value)
   {
     Values values;
     if (value->type() == Var)
@@ -558,7 +541,7 @@ namespace rego
     return values;
   }
 
-  Values Unifier::resolve_var(const Node& node)
+  Values UnifierDef::resolve_var(const Node& node)
   {
     Values values = check_with(node);
     if (!values.empty())
@@ -646,7 +629,7 @@ namespace rego
     return ValueDef::filter_by_rank(values);
   }
 
-  Values Unifier::call_function(
+  Values UnifierDef::call_function(
     const Location& var, const std::string& func_name, const Values& args)
   {
     Values values;
@@ -832,7 +815,7 @@ namespace rego
     return values;
   }
 
-  Values Unifier::enumerate(const Location& var, const Node& container_var)
+  Values UnifierDef::enumerate(const Location& var, const Node& container_var)
   {
     Values items;
     Values container_values = resolve_var(container_var);
@@ -889,14 +872,14 @@ namespace rego
     return items;
   }
 
-  std::string Unifier::str() const
+  std::string UnifierDef::str() const
   {
     std::ostringstream buf;
-    buf << *this;
+    buf << this;
     return buf.str();
   }
 
-  std::string Unifier::dependency_str() const
+  std::string UnifierDef::dependency_str() const
   {
     std::ostringstream buf;
     for (auto& [name, var] : m_variables)
@@ -914,7 +897,7 @@ namespace rego
     return buf.str();
   }
 
-  std::size_t Unifier::compute_dependency_score(const Node& unifyexpr)
+  std::size_t UnifierDef::compute_dependency_score(const Node& unifyexpr)
   {
     if (unifyexpr->type() == Local)
     {
@@ -960,7 +943,7 @@ namespace rego
     return score;
   }
 
-  bool Unifier::push_rule(const Location& rule)
+  bool UnifierDef::push_rule(const Location& rule)
   {
     LOG("Pushing rule: ", rule.view());
     LOG("Call stack: ", *m_call_stack);
@@ -976,7 +959,7 @@ namespace rego
     return false;
   }
 
-  void Unifier::pop_rule(const Location& rule)
+  void UnifierDef::pop_rule(const Location& rule)
   {
     if (m_call_stack->size() == 0)
     {
@@ -993,7 +976,7 @@ namespace rego
     m_call_stack->pop_back();
   }
 
-  Values Unifier::check_with(const Node& var)
+  Values UnifierDef::check_with(const Node& var)
   {
     std::string key_str = std::string(var->location().view());
     for (auto it = m_with_stack->rbegin(); it != m_with_stack->rend(); ++it)
@@ -1027,7 +1010,7 @@ namespace rego
     return {};
   }
 
-  Values Unifier::resolve_skip(const Node& skip)
+  Values UnifierDef::resolve_skip(const Node& skip)
   {
     assert(skip->type() == Skip);
     LOG("Resolving skip: ", to_json(skip->front()));
@@ -1122,7 +1105,7 @@ namespace rego
     return values;
   }
 
-  std::optional<RankedNode> Unifier::resolve_rulecomp(const Node& rulecomp)
+  std::optional<RankedNode> UnifierDef::resolve_rulecomp(const Node& rulecomp)
   {
     if (rulecomp->type() == DefaultRule)
     {
@@ -1147,7 +1130,7 @@ namespace rego
     {
       try
       {
-        body_result = rule_unifier(rulename, rulebody).unify();
+        body_result = rule_unifier(rulename, rulebody)->unify();
       }
       catch (const std::exception& e)
       {
@@ -1170,8 +1153,8 @@ namespace rego
         try
         {
           Unifier unifier = rule_unifier(rulename, value);
-          unifier.unify();
-          auto bindings = unifier.bindings();
+          unifier->unify();
+          auto bindings = unifier->bindings();
           Node result;
           for (auto& binding : bindings)
           {
@@ -1200,7 +1183,7 @@ namespace rego
     return std::nullopt;
   }
 
-  std::optional<RankedNode> Unifier::resolve_rulefunc(
+  std::optional<RankedNode> UnifierDef::resolve_rulefunc(
     const Node& rulefunc, const Nodes& args)
   {
     assert(rulefunc->type() == RuleFunc);
@@ -1224,7 +1207,7 @@ namespace rego
 
     try
     {
-      body_result = rule_unifier(rulename, rulebody).unify();
+      body_result = rule_unifier(rulename, rulebody)->unify();
     }
     catch (const std::exception& e)
     {
@@ -1252,8 +1235,8 @@ namespace rego
       try
       {
         Unifier unifier = rule_unifier(rulename, value);
-        unifier.unify();
-        auto bindings = unifier.bindings();
+        unifier->unify();
+        auto bindings = unifier->bindings();
         for (auto& binding : bindings)
         {
           Node var = binding / Var;
@@ -1273,7 +1256,7 @@ namespace rego
     return std::make_pair(index, value);
   }
 
-  std::optional<Node> Unifier::resolve_ruleset(const Nodes& ruleset)
+  std::optional<Node> UnifierDef::resolve_ruleset(const Nodes& ruleset)
   {
     Node argseq = NodeDef::create(ArgSeq);
 
@@ -1294,7 +1277,7 @@ namespace rego
       {
         try
         {
-          body_result = rule_unifier(rulename, rulebody).unify();
+          body_result = rule_unifier(rulename, rulebody)->unify();
         }
         catch (const std::exception& e)
         {
@@ -1317,8 +1300,8 @@ namespace rego
           try
           {
             Unifier unifier = rule_unifier(rulename, value);
-            unifier.unify();
-            auto bindings = unifier.bindings();
+            unifier->unify();
+            auto bindings = unifier->bindings();
             Node result;
             for (auto& binding : bindings)
             {
@@ -1354,7 +1337,7 @@ namespace rego
     return Resolver::set(argseq);
   }
 
-  std::optional<Node> Unifier::resolve_ruleobj(const Nodes& ruleobj)
+  std::optional<Node> UnifierDef::resolve_ruleobj(const Nodes& ruleobj)
   {
     Node argseq = NodeDef::create(ArgSeq);
 
@@ -1375,7 +1358,7 @@ namespace rego
       {
         try
         {
-          body_result = rule_unifier(rulename, rulebody).unify();
+          body_result = rule_unifier(rulename, rulebody)->unify();
         }
         catch (const std::exception& e)
         {
@@ -1398,8 +1381,8 @@ namespace rego
           try
           {
             Unifier unifier = rule_unifier(rulename, value);
-            unifier.unify();
-            auto bindings = unifier.bindings();
+            unifier->unify();
+            auto bindings = unifier->bindings();
             Node result;
             for (auto& binding : bindings)
             {
@@ -1439,7 +1422,7 @@ namespace rego
     return Resolver::object(argseq);
   }
 
-  Nodes Unifier::expressions() const
+  Nodes UnifierDef::expressions() const
   {
     Nodes terms;
     for (auto& [loc, var] : m_variables)
@@ -1452,7 +1435,7 @@ namespace rego
     return terms;
   }
 
-  Nodes Unifier::bindings() const
+  Nodes UnifierDef::bindings() const
   {
     Nodes bindings;
     for (auto& [loc, var] : m_variables)
@@ -1465,7 +1448,7 @@ namespace rego
     return bindings;
   }
 
-  void Unifier::push_with(const Node& withseq)
+  void UnifierDef::push_with(const Node& withseq)
   {
     LOG("pushing with lookup");
     ValuesLookup lookup;
@@ -1482,32 +1465,42 @@ namespace rego
     m_with_stack->push_back(lookup);
   }
 
-  void Unifier::pop_with()
+  void UnifierDef::pop_with()
   {
     LOG("popping with lookup");
     m_with_stack->pop_back();
   }
 
-  std::shared_ptr<Unifier> Unifier::nested(
-    const Location& rule, const Node& rulebody)
+  Unifier UnifierDef::create(
+    const Location& rule,
+    const Node& rulebody,
+    CallStack call_stack,
+    WithStack with_stack,
+    const BuiltIns& builtins,
+    UnifierCache cache)
   {
-    return std::shared_ptr<Unifier>(new Unifier(
-      rule, rulebody, m_call_stack, m_with_stack, m_builtins, 1, this));
+    if (cache->contains(rulebody))
+    {
+      Unifier unifier = cache->at(rulebody);
+      unifier->reset();
+      return unifier;
+    }
+    else
+    {
+      Unifier unifier = std::make_shared<UnifierDef>(
+        rule, rulebody, call_stack, with_stack, builtins, cache);
+      cache->insert({rulebody, unifier});
+      return unifier;
+    }
   }
 
-  Unifier Unifier::rule_unifier(const Location& rule, const Node& rulebody)
+  Unifier UnifierDef::rule_unifier(const Location& rule, const Node& rulebody)
   {
-    return Unifier(
-      rule,
-      rulebody,
-      m_call_stack,
-      m_with_stack,
-      m_builtins,
-      m_max_passes,
-      nullptr);
+    return create(
+      rule, rulebody, m_call_stack, m_with_stack, m_builtins, m_cache);
   }
 
-  Node Unifier::resolve_every(const Node& varseq, const Node& nestedbody)
+  Node UnifierDef::resolve_every(const Node& varseq, const Node& nestedbody)
   {
     Node item = varseq->at(0);
     Node key, val;
@@ -1547,7 +1540,7 @@ namespace rego
       {
         key->back() = item_node->at(0);
       }
-      Node result = rule_unifier(name, body).unify();
+      Node result = rule_unifier(name, body)->unify();
       if (result->type() == JSONFalse)
       {
         return JSONFalse;
@@ -1562,38 +1555,28 @@ namespace rego
     return JSONTrue;
   }
 
-  bool Unifier::is_variable(const Location& name) const
+  bool UnifierDef::is_variable(const Location& name) const
   {
     if (m_variables.contains(name))
     {
       return true;
     }
 
-    if (m_outer != nullptr)
-    {
-      return m_outer->is_variable(name);
-    }
-
     return false;
   }
 
-  Variable& Unifier::get_variable(const Location& name)
+  Variable& UnifierDef::get_variable(const Location& name)
   {
     if (m_variables.contains(name))
     {
       return m_variables.at(name);
     }
 
-    if (m_outer != nullptr)
-    {
-      return m_outer->get_variable(name);
-    }
-
     std::string name_str = std::string(name.view());
     throw std::runtime_error("Variable " + name_str + " not found");
   }
 
-  Values Unifier::resolve_compr(const Location& var, const Node& compr)
+  Values UnifierDef::resolve_compr(const Location& var, const Node& compr)
   {
     Values values;
     if (compr->type() == ArrayCompr)
diff --git a/templates/vscode/launch.json b/templates/vscode/launch.json
index 5b98cc15..889d85fb 100644
--- a/templates/vscode/launch.json
+++ b/templates/vscode/launch.json
@@ -254,6 +254,40 @@
       ],
       "cwd": "${workspaceFolder}/build/dist/"
     },    
+    {
+      "type": "lldb",
+      "request": "launch",
+      "name": "setops",
+      "program": "${workspaceFolder}/build/dist/bin/rego_interpreter",
+      "args": [
+        "-d",
+        "./examples/setops.rego",
+        "-q",
+        "data.setops.output"
+      ],
+      "cwd": "${workspaceFolder}/build/dist/"
+    },
+    {
+      "type": "lldb",
+      "request": "launch",
+      "name": "test",
+      "program": "${workspaceFolder}/build/dist/bin/rego_test",
+      "args": [
+        "tests/regocpp.yaml"
+      ],
+      "cwd": "${workspaceFolder}/build/dist/"
+    },
+    {
+      "type": "lldb",
+      "request": "launch",
+      "name": "yaml_test",
+      "program": "${workspaceFolder}/build/dist/bin/yaml_test",
+      "args": [
+        "test",
+        "-f"
+      ],
+      "cwd": "${workspaceFolder}/build/dist/"
+    },
     {
       "type": "lldb",
       "request": "launch",
diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt
index 65df245b..0e986939 100644
--- a/tools/CMakeLists.txt
+++ b/tools/CMakeLists.txt
@@ -10,47 +10,5 @@ target_link_libraries(rego_trieste
   PRIVATE
   regocpp::rego)
 
-set(tests_list
-scalars "-d\texamples/scalars.rego\t-q\tdata.scalars.greeting" "\"Hello\""
-objects "-d\texamples/objects.rego\t-q\tdata.objects.sites[1].name" "\"smoke1\""
-input_data "-d\texamples/data0.json\texamples/data1.json\texamples/objects.rego\t-i\texamples/input0.json\t-q\t[data.one, input.b, data.objects.sites[1]]" [=[\[{"bar": "Foo", "baz": 5, "be": true, "bop": 23.4}, "20", {"name": "smoke1"}\]]=]
-expressions "-d\texamples/expressions.rego\t-q\tdata.expressions.k > 27" "true"
-bodies "-d\texamples/bodies.rego\t-i\texamples/input0.json\t-q\tdata.bodies.e" [=[{"one": 15, "two": 15}]=]
-bodies_undef "-d\texamples/bodies.rego\t-i\texamples/input0.json\t-q\tdata.bodies.g" "undefined"
-sets "-d\texamples/sets.rego\t-q\tdata.sets.string_set" [=[\["four", "one", "two"\]]=]
-sets_membership "-d\texamples/sets.rego\t-q\tdata.sets.m" "true"
-var_keys "-d\texamples/objects.rego\t-q\tdata.objects.f" "false"
-multi0 "-d\texamples/multi.rego\t-i\texamples/input0.json\t-q\tdata.multi.a" "1"
-multi1 "-d\texamples/multi.rego\t-i\texamples/input1.json\t-q\tdata.multi.a" "41"
-multi2 "-d\texamples/multi.rego\t-i\texamples/input2.json\t-q\tdata.multi.a" "70"
-functions_sum "-d\texamples/functions.rego\t-i\texamples/input1.json\t-q\t[data.functions.f(1, 2, 3), data.functions.f(1, -1, 3)]" [=[\[6, 4\]]=]
-functions_prod "-d\texamples/functions.rego\t-i\texamples/input2.json\t-q\t[data.functions.f(1, 2, 3), data.functions.f(1, -1, 3)]" [=[\[6, -3\]]=]
-functions_ref "-d\texamples/functions.rego\t-q\t[data.functions.h(1, 2).a, data.functions.h(3, 4).b]" [=[\[1, 4\]]=]
-unify "-d\texamples/unify.rego\t-q\tdata.unify.output" [=[\[{"bar": 2, "foo": 1}, 3, true, true, "bar", 4, 3, \"foo\"\]]=]
-compare "-d\texamples/compare.rego\t-q\tdata.compare.output" [=[\[true, false, false, true\]]=]
-keywords "-d\texamples/keywords.rego\texamples/data2.json\t-i\texamples/input0.json\t-q\tdata.keywords.output" [=[\[4, "found", \["0", "a"\], {"a": 3, "b": 5, "c": true, "sum": 8}, true\]]=]
-imports "-d\texamples/scalars.rego\texamples/import0.rego\texamples/import1.rego\texamples/import2.rego\t-q\tdata.scalars.output" [=[\[305, true, 469\]]=]
-with "-d\texamples/bodies.rego\texamples/with.rego\t-q\tdata.bodies.output" [=[\[8, 2, \[\[100, 300\], {"bar": 300, "foo": 200}\], 3, 0\]]=]
-every "-d\texamples/every.rego\texamples/data2.json\t-q\tdata.every.output" "true"
-compr "-d\texamples/compr.rego\texamples/data2.json\t-q\tdata.compr.output" "true"
-setops "-d\texamples/setops.rego\t-q\tdata.setops.output" "true"
-)
-
-add_test(NAME rego_trieste COMMAND rego_trieste test -f WORKING_DIRECTORY $<TARGET_FILE_DIR:rego_trieste>)
-add_test(NAME query_math COMMAND rego_interpreter "-q" "x=5;x + (2 - 4 * 0.25) * -3 + 7.4")
-set_tests_properties(query_math PROPERTIES PASS_REGULAR_EXPRESSION "x = 5[\r\n\t ]*9.4")
-
-list(LENGTH tests_list tests_length)
-while(${tests_length} GREATER 0)
-  list(POP_FRONT tests_list test_name test_args test_regex)
-  string(REPLACE "\t" ";" test_args_list ${test_args})
-  add_test(NAME ${test_name}
-           COMMAND rego_interpreter ${test_args_list} -a debug
-           WORKING_DIRECTORY $<TARGET_FILE_DIR:rego_interpreter>)
-  set_tests_properties(${test_name} PROPERTIES
-           PASS_REGULAR_EXPRESSION "^${test_regex}")
-  list(LENGTH tests_list tests_length)
-endwhile()
-
 add_custom_command(TARGET rego_interpreter POST_BUILD
                    COMMAND ${CMAKE_COMMAND} -E copy_directory ${CMAKE_CURRENT_SOURCE_DIR}/examples $<TARGET_FILE_DIR:rego_interpreter>/examples)
diff --git a/tools/driver.cc b/tools/driver.cc
index 89ab11d1..1c1bfc5f 100644
--- a/tools/driver.cc
+++ b/tools/driver.cc
@@ -1,4 +1,4 @@
-#include "lang.h"
+#include "rego.h"
 
 int main(int argc, char** argv)
 {
diff --git a/tools/main.cc b/tools/main.cc
index 123565ca..712053fa 100644
--- a/tools/main.cc
+++ b/tools/main.cc
@@ -1,5 +1,4 @@
-#include "interpreter.h"
-#include "log.h"
+#include "rego.h"
 
 int main(int argc, char** argv)
 {
@@ -7,6 +6,9 @@ int main(int argc, char** argv)
 
   app.set_help_all_flag("--help-all", "Expand all help");
 
+  std::string query_expr;
+  app.add_option("query,-q,--query", query_expr, "Query")->required();
+
   std::vector<std::filesystem::path> data_paths;
   app.add_option("-d,--data", data_paths, "Data/Policy files");
 
@@ -16,9 +18,6 @@ int main(int argc, char** argv)
   std::filesystem::path output;
   app.add_option("-a,--ast", output, "Output the AST");
 
-  std::string query_expr;
-  app.add_option("-q,--query", query_expr, "Query")->required();
-
   bool enable_logging{false};
   app.add_flag("-l,--logging", enable_logging, "Enable logging");
 


Test patch:
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
new file mode 100644
index 00000000..8033ff10
--- /dev/null
+++ b/tests/CMakeLists.txt
@@ -0,0 +1,23 @@
+set( SOURCES
+test_case.cc
+parse.cc
+rego_test.cc
+)
+
+add_executable(rego_test main.cc ${SOURCES})
+target_link_libraries(rego_test
+  PRIVATE 
+  regocpp::rego)
+
+add_executable(yaml_trieste driver.cc ${SOURCES})
+
+target_link_libraries(yaml_trieste
+  PRIVATE
+  regocpp::rego)
+
+add_test(NAME rego_trieste COMMAND rego_trieste test -f WORKING_DIRECTORY $<TARGET_FILE_DIR:rego_trieste>)
+add_test(NAME yaml_trieste COMMAND yaml_trieste test -f WORKING_DIRECTORY $<TARGET_FILE_DIR:yaml_trieste>)
+add_test(NAME rego_test COMMAND rego_test regocpp.yaml WORKING_DIRECTORY $<TARGET_FILE_DIR:rego_test>)
+
+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)
diff --git a/tests/README.md b/tests/README.md
new file mode 100644
index 00000000..de1140a1
--- /dev/null
+++ b/tests/README.md
@@ -0,0 +1,3 @@
+start
+- launching place for tokens etc.
+- if spaces, count them
\ No newline at end of file
diff --git a/tests/driver.cc b/tests/driver.cc
new file mode 100644
index 00000000..6bea9292
--- /dev/null
+++ b/tests/driver.cc
@@ -0,0 +1,6 @@
+#include "rego_test.h"
+
+int main(int argc, char** argv)
+{
+  return rego_test::driver().run(argc, argv);
+}
diff --git a/tests/main.cc b/tests/main.cc
new file mode 100644
index 00000000..a99d78c9
--- /dev/null
+++ b/tests/main.cc
@@ -0,0 +1,87 @@
+#include "test_case.h"
+
+const std::string Green = "\x1b[32m";
+const std::string Reset = "\x1b[0m";
+const std::string Red = "\x1b[31m";
+
+int main(int argc, char** argv)
+{
+  CLI::App app;
+
+  app.set_help_all_flag("--help-all", "Expand all help");
+
+  std::vector<std::filesystem::path> case_paths;
+  app.add_option("case,-c,--case", case_paths, "Test case YAML files");
+
+  std::filesystem::path debug_path;
+  app.add_option(
+    "-a,--ast", debug_path, "Output the AST (debugging for test case parser)");
+
+  try
+  {
+    app.parse(argc, argv);
+  }
+  catch (const CLI::ParseError& e)
+  {
+    return app.exit(e);
+  }
+
+  int failures = 0;
+  for (auto path : case_paths)
+  {
+    std::vector<rego_test::TestCase> test_cases;
+    auto maybe_test_cases = rego_test::TestCase::load(path, debug_path);
+    if (maybe_test_cases.has_value())
+    {
+      test_cases = *maybe_test_cases;
+    }
+    else
+    {
+      std::cout << "Unable to parse test cases in " << path << std::endl;
+      continue;
+    }
+
+    std::cout << path << std::endl;
+
+    for (auto& test_case : test_cases)
+    {
+      try
+      {
+        auto start = std::chrono::steady_clock::now();
+        auto result = test_case.run(argv[0]);
+        auto end = std::chrono::steady_clock::now();
+        const std::chrono::duration<double> elapsed = end - start;
+
+        if (result.passed)
+        {
+          std::cout << Green << "  PASS: " << Reset << test_case.note()
+                    << std::fixed << std::setw(62 - test_case.note().length())
+                    << std::internal << std::setprecision(2) << elapsed.count()
+                    << " sec" << std::endl;
+        }
+        else
+        {
+          failures++;
+          std::cout << std::setw(70) << Red << "  FAIL: " << Reset
+                    << test_case.note() << std::fixed
+                    << std::setw(62 - test_case.note().length())
+                    << std::internal << std::setprecision(2) << elapsed.count()
+                    << " sec" << std::endl;
+          std::cout << "  " << result.error << std::endl;
+        }
+      }
+      catch (const std::exception& e)
+      {
+        failures++;
+        std::cout << std::setw(70) << Red << "  FAIL: " << Reset
+                  << test_case.note() << std::endl;
+        std::cout << "  " << e.what() << std::endl;
+        continue;
+      }
+    }
+
+    std::cout << std::endl;
+  }
+
+  return failures;
+}
\ No newline at end of file
diff --git a/tests/parse.cc b/tests/parse.cc
new file mode 100644
index 00000000..caf30743
--- /dev/null
+++ b/tests/parse.cc
@@ -0,0 +1,263 @@
+#include "rego_test.h"
+
+namespace rego_test
+{
+  const std::string alpha =
+    "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
+  const std::string numeric = "0123456789";
+  const std::string alphanumeric = alpha + numeric;
+  const std::size_t max_string_length = 10;
+
+  template <typename T>
+  std::string rand_string(T& rnd)
+  {
+    // TODO add valid non-alphanum characters
+    std::ostringstream buf;
+    std::size_t length = rnd() % max_string_length;
+    for (std::size_t i = 0; i < length; ++i)
+    {
+      buf << alphanumeric[rnd() % alphanumeric.size()];
+    }
+    return buf.str();
+  }
+
+  struct Block
+  {
+    std::size_t indent;
+  };
+
+  Parse parser()
+  {
+    Parse p(depth::file);
+    auto indents = std::make_shared<std::vector<std::size_t>>();
+    indents->push_back(0);
+    auto indent = std::make_shared<std::size_t>(0);
+    auto string_indent = std::make_shared<std::size_t>(0);
+
+    p("start",
+      {
+        // Line comment.
+        "#[^\n]*" >> [](auto&) {},
+
+        // end of file terminates
+        "\r*\n$" >>
+          [indents](auto& m) {
+            while (!indents->empty())
+            {
+              m.term({Block});
+              indents->pop_back();
+            }
+          },
+
+        "\r*\n" >>
+          [indent](auto& m) {
+            m.term();
+            *indent = 0;
+            m.mode("indent");
+          },
+
+        // Brace.
+        R"((\{)[[:blank:]]*)" >> [](auto& m) { m.push(Brace, 1); },
+
+        R"(\})" >>
+          [](auto& m) {
+            m.term();
+            m.pop(Brace);
+          },
+
+        // Square.
+        R"((\[)[[:blank:]]*)" >> [](auto& m) { m.push(Square, 1); },
+
+        R"(\])" >>
+          [](auto& m) {
+            m.term();
+            m.pop(Square);
+          },
+
+        // Comma.
+        ",[[:blank:]]*" >> [](auto& m) { m.term(); },
+
+        // Double quote string
+        R"re2("([^"]*)")re2" >> [](auto& m) { m.add(String, 1); },
+
+        // Single quote string
+        R"('([^'^\\]*)')" >> [](auto& m) { m.add(String, 1); },
+
+        // KeyValue.
+        ":[ ]*" >> [](auto& m) { m.add(Colon); },
+
+        // Literal string
+        "\\|" >> [](auto& m) { m.push(LiteralString); },
+
+        // Folded string
+        ">" >> [](auto& m) { m.push(FoldedString); },
+
+        // Entry.
+        "-[ ]*" >> [](auto& m) { m.add(Hyphen); },
+
+        // True.
+        "true\\b" >> [](auto& m) { m.add(True); },
+
+        // False.
+        "false\\b" >> [](auto& m) { m.add(False); },
+
+        // Null.
+        "null\\b" >> [](auto& m) { m.add(Null); },
+
+        // String
+        R"([a-zA-Z_][a-zA-Z0-9_!\/ \.=\+\-\*\[\]\(\);><\|]*)" >>
+          [](auto& m) { m.extend(String); },
+
+        // Float.
+        R"([[:digit:]]+\.[[:digit:]]+(?:e[+-]?[[:digit:]]+)?\b)" >>
+          [](auto& m) { m.add(Float); },
+
+        // Int.
+        R"([[:digit:]]+\b)" >> [](auto& m) { m.add(Integer); },
+      });
+
+    p("indent",
+      {// end of file terminates
+       "\r*\n$" >>
+         [indents](auto& m) {
+           while (!indents->empty())
+           {
+             m.term({Block});
+             indents->pop_back();
+           }
+         },
+
+       "\r*\n" >>
+         [indent, string_indent](auto& m) {
+           if (*indent == *string_indent)
+           {
+             m.add(Blank);
+             m.term();
+           }
+           *indent = 0;
+         },
+
+       // Space
+       " " >>
+         [indent, string_indent](auto& m) {
+           if (*string_indent > 0 && *indent == *string_indent)
+           {
+             m.mode("multiline-string");
+             m.add(String);
+           }
+           else
+           {
+             *indent = *indent + 1;
+           }
+         },
+
+       // Hyphen
+       "-[ ]*" >>
+         [indent, string_indent, indents](auto& m) {
+           m.mode("start");
+           if (m.in(LiteralString) || m.in(FoldedString))
+           {
+             *string_indent = 0;
+             if (m.in(LiteralString))
+             {
+               m.pop(LiteralString);
+             }
+             else
+             {
+               m.pop(FoldedString);
+             }
+           }
+
+           while (*indent < indents->back())
+           {
+             m.term({Block});
+             indents->pop_back();
+           }
+
+           *indent += 2;
+           m.add(Hyphen);
+           m.push(Block);
+           indents->push_back(*indent);
+         },
+
+       // Character
+       "." >>
+         [indent, string_indent, indents](auto& m) {
+           std::string mode = "start";
+           if (m.in(LiteralString) || m.in(FoldedString))
+           {
+             mode = "multiline-string";
+             if (*string_indent == 0)
+             {
+               *string_indent = *indent;
+             }
+             else if (*indent < *string_indent)
+             {
+               mode = "start";
+               *string_indent = 0;
+               if (m.in(LiteralString))
+               {
+                 m.pop(LiteralString);
+               }
+               else
+               {
+                 m.pop(FoldedString);
+               }
+             }
+           }
+           else
+           {
+             if (*indent > indents->back())
+             {
+               m.push(Block);
+               indents->push_back(*indent);
+             }
+           }
+
+           while (*indent < indents->back())
+           {
+             m.term({Block});
+             indents->pop_back();
+           }
+           m.term();
+           m.mode(mode);
+           m.add(String);
+         }});
+
+    p("multiline-string",
+      {// end of file terminates
+       "\r*\n$" >>
+         [indents](auto& m) {
+           while (!indents->empty())
+           {
+             m.term({Block});
+             indents->pop_back();
+           }
+         },
+
+       "\r*\n" >>
+         [indent](auto& m) {
+           *indent = 0;
+           m.term();
+           m.mode("indent");
+         },
+
+       // Character
+       "." >> [](auto& m) { m.extend(String); }});
+
+    p.gen({
+      Integer >> [](auto& rnd) { return std::to_string(rnd() % 100); },
+      Float >>
+        [](auto& rnd) {
+          std::uniform_real_distribution<> dist(-10.0, 10.0);
+          return std::to_string(dist(rnd));
+        },
+      True >> [](auto&) { return "true"; },
+      False >> [](auto&) { return "false"; },
+      Null >> [](auto&) { return "null"; },
+      String >> [](auto& rnd) { return rand_string(rnd); },
+    });
+
+    return p;
+  }
+}
\ No newline at end of file
diff --git a/tests/rego_test.cc b/tests/rego_test.cc
new file mode 100644
index 00000000..7cb146e9
--- /dev/null
+++ b/tests/rego_test.cc
@@ -0,0 +1,369 @@
+#include "rego_test.h"
+
+namespace rego_test
+{
+  const auto inline SequenceToken = T(Sequence) / T(Mapping) / T(Scalar);
+  const auto inline ScalarToken =
+    T(String) / T(Integer) / T(Float) / T(True) / T(False) / T(Null);
+  const auto inline ValToken = T(Mapping) / T(Sequence) / T(Scalar);
+
+  PassDef entry()
+  {
+    return {
+      In(Top) * T(File)[File] >> [](Match& _) { return Block << *_[File]; },
+
+      In(Group) * T(Brace)[Brace] >>
+        [](Match& _) { return Block << *_[Brace]; },
+
+      In(Group) * T(Square)[Square] >>
+        [](Match& _) {
+          Node seq = NodeDef::create(Seq);
+          for (auto group : *_(Square))
+          {
+            seq << (Entry << group);
+          }
+          return seq;
+        },
+
+      In(Group) * (T(Hyphen) * (T(Block) / T(Group))[Entry]) >>
+        [](Match& _) { return Entry << _(Entry); },
+
+      In(Group) * (T(Hyphen)[Hyphen] * T(Integer)[Integer]) >>
+        [](Match& _) {
+          Location loc0 = _(Hyphen)->location();
+          Location loc1 = _(Integer)->location();
+          Location loc(loc0.source, loc0.pos, loc1.pos + loc1.len - loc0.pos);
+          return Integer ^ loc;
+        },
+
+      In(Group) * (T(Hyphen)[Hyphen] * T(Float)[Float]) >>
+        [](Match& _) {
+          Location loc0 = _(Hyphen)->location();
+          Location loc1 = _(Float)->location();
+          Location loc(loc0.source, loc0.pos, loc1.pos + loc1.len - loc0.pos);
+          return Float ^ loc;
+        },
+
+      (In(LiteralString) / In(FoldedString)) *
+          (T(Group) << T(String)[String]) >>
+        [](Match& _) { return _(String); },
+
+      (In(LiteralString) / In(FoldedString)) * (T(Group) << T(Blank)) >>
+        [](Match&) { return String ^ ""; },
+
+      // errors
+      In(Group) * T(Hyphen)[Hyphen] >>
+        [](Match& _) {
+          return err(_(Hyphen), "Invalid sequence entry declaration");
+        },
+
+      In(Group) * T(Blank)[Blank] >>
+        [](Match& _) { return err(_(Blank), "Invalid blank line"); },
+
+      In(LiteralString) * T(Group)[Group] >>
+        [](Match& _) {
+          return err(_(Group), "Invalid literal string element");
+        },
+
+      In(FoldedString) * T(Group)[Group] >>
+        [](Match& _) { return err(_(Group), "Invalid folded string element"); },
+
+    };
+  }
+
+  PassDef sequence()
+  {
+    return {
+      In(Group) * (T(Entry)[Head] * T(Entry)++[Tail] * End) >>
+        [](Match& _) { return Sequence << _(Head) << _[Tail]; },
+
+      In(Group) * (T(LiteralString) << (T(String)++[String])) >>
+        [](Match& _) {
+          std::ostringstream buf;
+          std::string sep = "";
+          for (auto it = _[String].first; it != _[String].second; ++it)
+          {
+            buf << sep << (*it)->location().view();
+            sep = "\n";
+          }
+          return String ^ buf.str();
+        },
+
+      In(Group) * (T(FoldedString) << (T(String)++[String])) >>
+        [](Match& _) {
+          std::ostringstream buf;
+          std::string sep = "";
+          for (auto it = _[String].first; it != _[String].second; ++it)
+          {
+            buf << sep << (*it)->location().view();
+            sep = " ";
+          }
+          return String ^ buf.str();
+        },
+
+      // errors
+
+      In(LiteralString) * T(Group)[Group] >>
+        [](Match& _) {
+          return err(_(Group), "Invalid literal string element");
+        },
+
+      In(FoldedString) * T(Group)[Group] >>
+        [](Match& _) { return err(_(Group), "Invalid folded string element"); },
+
+      In(Group) * T(Entry)[Entry] >>
+        [](Match& _) { return err(_(Entry), "Invalid entry location"); }};
+  }
+
+  PassDef scalar()
+  {
+    return {
+      In(Group) * ScalarToken[Scalar] >>
+        [](Match& _) { return Scalar << _(Scalar); },
+    };
+  }
+
+  PassDef keyvalue()
+  {
+    return {
+      (In(Entry) / In(Block)) *
+          (T(Group) << ((T(Scalar) << T(String))[Key] * T(Colon) * Any[Val])) >>
+        [](Match& _) {
+          Location key_loc = _(Key)->location();
+          std::string key = std::string(key_loc.view());
+          std::size_t length = key.find(' ');
+          if (length != std::string::npos)
+          {
+            key_loc.len = length;
+          }
+
+          return KeyValue << (Key ^ key_loc) << (Group << _(Val));
+        },
+
+      (In(Entry) / In(Block)) *
+          ((T(Group) << ((T(Scalar) << T(String))[Key] * T(Colon) * End)) *
+           T(Group)[Val]) >>
+        [](Match& _) {
+          Location key_loc = _(Key)->location();
+          std::string key = std::string(key_loc.view());
+          std::size_t length = key.find(' ');
+          if (length != std::string::npos)
+          {
+            key_loc.len = length;
+          }
+
+          return KeyValue << (Key ^ key_loc) << _(Val);
+        },
+
+      In(Entry) * T(KeyValue)[KeyValue] >>
+        [](Match& _) { return Group << (Block << _(KeyValue)); },
+
+      In(Entry) * (T(Block) << (T(Group)[Group] * End)) >>
+        [](Match& _) { return _(Group); },
+
+      In(Entry) * T(Block)[Block] >> [](Match& _) { return Group << _(Block); },
+
+      // errors
+      In(Group) * T(Colon)[Colon] >>
+        [](Match& _) { return err(_(Colon), "Invalid key/value declaration"); },
+
+      In(Group) * (Any * Any[Val]) >>
+        [](Match& _) { return err(_(Val), "Syntax error"); },
+    };
+  }
+
+  PassDef mapping()
+  {
+    return {
+      In(Group) * (T(Block) << (T(KeyValue)++[Mapping] * End)) >>
+        [](Match& _) { return Mapping << _[Mapping]; },
+
+      // errors
+      (In(Entry) / In(Group)) * T(Block)[Block] >>
+        [](Match& _) { return err(_(Block), "Invalid indented block"); },
+
+      In(Group) * T(KeyValue)[KeyValue] >>
+        [](Match& _) {
+          return err(_(KeyValue), "Invalid key/value declaration");
+        },
+    };
+  }
+
+  PassDef structure()
+  {
+    return {
+      In(KeyValue) * (T(Group) << ValToken[Val]) >>
+        [](Match& _) { return _(Val); },
+
+      In(Entry) * (T(Group) << ValToken[Val]) >>
+        [](Match& _) { return _(Val); },
+
+      In(Top) * (T(Block) << (T(KeyValue)[KeyValue] * End)) >>
+        [](Match& _) { return Document << _(KeyValue); },
+
+      // errors
+      In(Top) * T(Block)[Block] >>
+        [](Match& _) { return err(_(Block), "Invalid rego test case"); },
+    };
+  }
+
+  bool name_equals(const Node& node, const std::set<std::string>& names)
+  {
+    std::string name = std::string(node->location().view());
+    return names.contains(name);
+  }
+
+  PassDef to_rego()
+  {
+    return {
+      In(Mapping) *
+          (T(KeyValue)
+           << ((T(Key)[Key]([](auto& n) {
+                 return name_equals(*n.first, {"data", "input"});
+               })) *
+               T(Mapping)[Mapping])) >>
+        [](Match& _) {
+          return KeyValue << _(Key)
+                          << (rego::File
+                              << (Group
+                                  << (rego::Brace
+                                      << (rego::List << *_[Mapping]))));
+        },
+
+      In(Mapping) *
+          (T(KeyValue)
+           << ((T(Key)[Key]([](auto& n) {
+                 return name_equals(*n.first, {"want_result"});
+               })) *
+               T(Mapping)[Mapping])) >>
+        [](Match& _) {
+          return KeyValue << _(Key) << (WantResult << *_[Mapping]);
+        },
+
+      In(rego::List) *
+          (T(KeyValue)
+           << (T(Key)[Key] * (T(Mapping) / T(Sequence) / T(Scalar))[Val])) >>
+        [](Match& _) { return Group << _(Key) << rego::Colon << _(Val); },
+
+      In(rego::List) * (T(Scalar) / T(Mapping) / T(Sequence))[Val] >>
+        [](Match& _) { return Group << _(Val); },
+
+      In(rego::List) * T(Entry)[Entry] >>
+        [](Match& _) { return _(Entry)->front(); },
+
+      In(Group) * T(Key)[Key] >>
+        [](Match& _) { return rego::JSONString ^ _(Key); },
+
+      (In(rego::Scalar) / In(Group)) * (T(Scalar) << T(String)[String]) >>
+        [](Match& _) { return rego::JSONString ^ _(String); },
+
+      (In(rego::Scalar) / In(Group)) * (T(Scalar) << T(Integer)[Integer]) >>
+        [](Match& _) { return rego::JSONInt ^ _(Integer); },
+
+      (In(rego::Scalar) / In(Group)) * (T(Scalar) << T(Float)[Float]) >>
+        [](Match& _) { return rego::JSONFloat ^ _(Float); },
+
+      (In(rego::Scalar) / In(Group)) * (T(Scalar) << T(True)[True]) >>
+        [](Match& _) { return rego::JSONTrue ^ _(True); },
+
+      (In(rego::Scalar) / In(Group)) * (T(Scalar) << T(False)[False]) >>
+        [](Match& _) { return rego::JSONFalse ^ _(False); },
+
+      (In(rego::Scalar) / In(Group)) * (T(Scalar) << T(Null)[Null]) >>
+        [](Match& _) { return rego::JSONNull ^ _(Null); },
+
+      In(Group) * T(Mapping)[Mapping] >>
+        [](Match& _) {
+          // test for set
+          return rego::Brace << (rego::List << *_[Mapping]);
+        },
+
+      In(Group) * T(Sequence)[Sequence] >>
+        [](Match& _) { return rego::Square << (rego::List << *_[Sequence]); },
+
+      In(WantResult) * (T(KeyValue) << (T(Key)[Key] * Any[Val])) >>
+        [](Match& _) {
+          return rego::Binding << (rego::Var ^ _(Key))
+                               << (rego::Term << _(Val));
+        },
+
+      In(rego::Term) * T(Mapping)[Mapping] >>
+        [](Match& _) {
+          bool is_set = true;
+          for (auto& keyvalue : *_(Mapping))
+          {
+            Node val = keyvalue / Val;
+            if (val->type() != Scalar)
+            {
+              is_set = false;
+              break;
+            }
+            val = val->front();
+            if (val->type() != Null)
+            {
+              is_set = false;
+              break;
+            }
+          }
+
+          if (is_set)
+          {
+            Node set = NodeDef::create(rego::Set);
+            for (auto& keyvalue : *_(Mapping))
+            {
+              set
+                << (rego::Term
+                    << (rego::Scalar << (rego::JSONString ^ keyvalue / Key)));
+            }
+            return set;
+          }
+
+          return rego::Object << *_[Mapping];
+        },
+
+      In(rego::Object) * (T(KeyValue) << (T(Key)[Key] * Any[Val])) >>
+        [](Match& _) {
+          return rego::ObjectItem << (rego::Key ^ _(Key))
+                                  << (rego::Term << _(Val));
+        },
+
+      In(rego::Term) * T(Scalar)[Scalar] >>
+        [](Match& _) { return rego::Scalar << _(Scalar); },
+
+      In(rego::Term) * T(Sequence)[Sequence] >>
+        [](Match& _) { return rego::Array << *_[Sequence]; },
+
+      In(rego::Array) * T(Entry)[Entry] >>
+        [](Match& _) { return rego::Term << *_[Entry]; }};
+  }
+
+  std::vector<PassCheck> passes()
+  {
+    return {
+      {"entry", entry(), &wf_pass_entry},
+      {"sequence", sequence(), &wf_pass_sequence},
+      {"scalar", scalar(), &wf_pass_scalar},
+      {"keyvalue", keyvalue(), &wf_pass_keyvalue},
+      {"mapping", mapping(), &wf_pass_mapping},
+      {"structure", structure(), &wf_pass_structure},
+      {"to_rego", to_rego(), &wf_pass_to_rego}};
+  }
+  Driver& driver()
+  {
+    static Driver d(
+      "rego_test",
+      nullptr,
+      parser(),
+      wf_parser,
+      {
+        {"entry", entry(), wf_pass_entry},
+        {"sequence", sequence(), wf_pass_sequence},
+        {"scalar", scalar(), wf_pass_scalar},
+        {"keyvalue", keyvalue(), wf_pass_keyvalue},
+        {"mapping", mapping(), wf_pass_mapping},
+        {"structure", structure(), wf_pass_structure},
+        {"to_rego", to_rego(), wf_pass_to_rego},
+      });
+    return d;
+  }
+}
\ No newline at end of file
diff --git a/tests/rego_test.h b/tests/rego_test.h
new file mode 100644
index 00000000..e5b59d6c
--- /dev/null
+++ b/tests/rego_test.h
@@ -0,0 +1,161 @@
+#pragma once
+
+// Trieste-based rewriter for Rego YAML test cases
+
+#include "rego.h"
+#include "trieste/driver.h"
+#include "trieste/token.h"
+#include "wf.h"
+
+namespace rego_test
+{
+  using namespace trieste;
+  using namespace wf::ops;
+
+  inline const auto Document = TokenDef("yaml-document", flag::symtab);
+  inline const auto Block = TokenDef("yaml-block");
+  inline const auto Sequence = TokenDef("yaml-sequence");
+  inline const auto Entry = TokenDef("yaml-entry");
+  inline const auto KeyValue = TokenDef("yaml-key-value", flag::lookdown);
+  inline const auto Mapping = TokenDef("yaml-mapping", flag::symtab);
+  inline const auto Key = TokenDef("yaml-key", flag::print);
+  inline const auto Val = TokenDef("yaml-val");
+  inline const auto String = TokenDef("yaml-string", flag::print);
+  inline const auto Integer = TokenDef("yaml-integer", flag::print);
+  inline const auto Float = TokenDef("yaml-float", flag::print);
+  inline const auto True = TokenDef("yaml-true", flag::print);
+  inline const auto False = TokenDef("yaml-false", flag::print);
+  inline const auto Null = TokenDef("yaml-null", flag::print);
+  inline const auto LiteralString = TokenDef("yaml-literal-string");
+  inline const auto FoldedString = TokenDef("yaml-folded-string");
+  inline const auto Scalar = TokenDef("yaml-scalar");
+  inline const auto Data = TokenDef("yaml-data");
+  inline const auto Input = TokenDef("yaml-input");
+  inline const auto Colon = TokenDef("yaml-:");
+  inline const auto Hyphen = TokenDef("yaml--");
+  inline const auto WantResult = TokenDef("yaml-want-result");
+  inline const auto Head = TokenDef("yaml-head");
+  inline const auto Tail = TokenDef("yaml-tail");
+  inline const auto Blank = TokenDef("yaml-blank");
+  inline const auto Brace = TokenDef("yaml-brace");
+  inline const auto Square = TokenDef("yaml-square");
+
+  inline const auto wf_parse_tokens = Block | String | Integer | Float | True |
+    False | Null | Colon | LiteralString | FoldedString | Hyphen | Blank |
+    Brace | Square;
+
+  // clang-format off
+  inline const auto wf_parser =
+    (Top <<= File)
+    | (File <<= Group++[1])
+    | (Block <<= Group++[1])
+    | (LiteralString <<= Group++[1])
+    | (FoldedString <<= Group++[1])
+    | (Brace <<= Group++[1])
+    | (Square <<= Group++[1])
+    | (Group <<= wf_parse_tokens++[1])
+    ;
+  // clang-format on
+
+  inline const auto wf_entry_tokens = Block | String | Integer | Float | True |
+    False | Null | Colon | Entry | LiteralString | FoldedString;
+
+  // clang-format off
+  inline const auto wf_pass_entry =
+    wf_parser
+    | (Top <<= Block)
+    | (LiteralString <<= String++[1])
+    | (FoldedString <<= String++[1])
+    | (Entry <<= Block | Group)
+    | (Group <<= wf_entry_tokens++[1])
+    ;
+  // clang-format on    
+
+  inline const auto wf_sequence_tokens =
+    Block | String | Integer | Float | True | False | Null | Colon | Sequence;
+
+  // clang-format off
+  inline const auto wf_pass_sequence =
+    wf_pass_entry
+    | (Sequence <<= Entry++[1])
+    | (Group <<= wf_sequence_tokens++[1])
+    ;
+  // clang-format on   
+
+  // clang-format off
+  inline const auto wf_pass_scalar =
+    wf_pass_sequence
+    | (Scalar <<= String | Integer | Float | True | False | Null)
+    | (Group <<= (Block | Scalar | Sequence | Colon)++[1])
+    ;
+  // clang-format on
+
+  // clang-format off
+  inline const auto wf_pass_keyvalue =
+    wf_pass_scalar
+    | (KeyValue <<= Key * Group)
+    | (Block <<= (Group | KeyValue)++[1])
+    | (Group <<= Block | Scalar | Sequence | KeyValue)
+    ;
+  // clang-format on
+
+  // clang-format off
+  inline const auto wf_pass_mapping =
+    wf_pass_keyvalue
+    | (Mapping <<= KeyValue++[1])
+    | (Entry <<= Group)
+    | (Group <<= Mapping | Scalar | Sequence)
+    ;
+
+  // clang-format off
+  inline const auto wf_pass_structure =
+    (Top <<= Document)
+    | (Document <<= KeyValue)
+    | (KeyValue <<= Key * (Val >>= Scalar | Sequence | Mapping))[Key]
+    | (Sequence <<= Entry++[1])
+    | (Mapping <<= KeyValue++[1])
+    | (Entry <<= Scalar | Sequence | Mapping)
+    | (Scalar <<= String | Integer | Float | True | False | Null)
+    ;
+  // clang-format on
+
+  const auto inline wf_rego_term_tokens =
+    rego::Scalar | rego::Object | rego::Array | rego::Set;
+  const auto inline wf_rego_scalar_tokens = rego::JSONInt | rego::JSONFloat |
+    rego::JSONString | rego::JSONTrue | rego::JSONFalse | rego::JSONNull;
+
+  // clang-format off
+  inline const auto wf_pass_to_rego =
+    wf_pass_structure
+    | (KeyValue <<= Key * (Val >>= Mapping | Sequence | Scalar | rego::File | WantResult))[Key]
+    | (rego::File <<= Group++)
+    | (rego::Brace <<= rego::List)
+    | (rego::Square <<= rego::List)
+    | (rego::List <<= Group++)
+    | (Group <<= rego::wf_parse_tokens++[1])
+    | (WantResult <<= rego::Binding++[1])
+    | (rego::Binding <<= rego::Var * rego::Term)
+    | (rego::Term <<= wf_rego_term_tokens)
+    | (rego::Object <<= rego::ObjectItem++[1])
+    | (rego::ObjectItem <<= rego::Key * rego::Term)
+    | (rego::Array <<= rego::Term++)
+    | (rego::Set <<= rego::Term++)
+    | (rego::Scalar <<= wf_rego_scalar_tokens)
+    ;
+  // clang-format on
+
+  inline auto err(Node node, const std::string& msg)
+  {
+    return Error << (ErrorMsg ^ msg) << (ErrorAst << node->clone());
+  }
+
+  inline auto err(const std::string& msg)
+  {
+    return Error << (ErrorMsg ^ msg);
+  }
+
+  Parse parser();
+  Driver& driver();
+  using PassCheck = std::tuple<std::string, Pass, const wf::Wellformed*>;
+  std::vector<PassCheck> passes();
+}
\ No newline at end of file
diff --git a/tests/regocpp.yaml b/tests/regocpp.yaml
new file mode 100644
index 00000000..5b72bfba
--- /dev/null
+++ b/tests/regocpp.yaml
@@ -0,0 +1,1063 @@
+cases:
+- modules:
+  - |
+    package scalars
+    
+    greeting := "Hello"
+    max_height := 42
+    pi := 3.14159
+    allowed := true
+    location := null
+  note       : regocpp/scalars
+  query      : data.scalars.greeting = x
+  want_result:
+    x: Hello
+- data:
+    one:
+      bar: Foo
+      baz: 5
+      be : true
+      bop: 23.4
+    two:
+      bar: Bar
+      baz: 12.3
+      be : false
+      bop: 42
+    three:
+      bar: Baz
+      baz: 15
+      be : true
+      bop: 4.23
+  input:
+    a: 10
+    b: "20"
+    c: 30.0
+    d: true
+  modules:
+  - |
+    package objects
+    
+    rect := {`width`: 2, "height": 4}
+    cube := {"width": 3, `height`: 4, "depth": 5}
+    a := 42
+    b := false
+    c := null
+    d := {"a": a, "x": [b, c]}
+    index := 1
+    shapes := [rect, cube]
+    names := ["prod", `smoke1`, "dev"]
+    sites := [{"name": "prod"}, {"name": names[index]}, {"name": "dev"}]
+    e := {
+        a: "foo",
+        "three": c,
+        names[2]: b,
+        "four": d,
+    }
+    f := e["dev"]
+  note : regocpp/input_data
+  query: "[data.one, input.b, data.objects.sites[1]] = x"
+  want_result:
+    x:
+    - bar: Foo
+      baz: 5
+      be: true
+      bop: 23.4
+    - "20"
+    - name: smoke1
+- modules:
+  - |
+    package objects
+    
+    rect := {`width`: 2, "height": 4}
+    cube := {"width": 3, `height`: 4, "depth": 5}
+    a := 42
+    b := false
+    c := null
+    d := {"a": a, "x": [b, c]}
+    index := 1
+    shapes := [rect, cube]
+    names := ["prod", `smoke1`, "dev"]
+    sites := [{"name": "prod"}, {"name": names[index]}, {"name": "dev"}]
+    e := {
+        a: "foo",
+        "three": c,
+        names[2]: b,
+        "four": d,
+    }
+    f := e["dev"]
+  note : regocpp/objects
+  query: data.objects.sites[1].name = x
+  want_result:
+    x: smoke1
+- note: regocpp/query_math
+  query: x=5;y=x + (2 - 4 * 0.25) * -3 + 7.4
+  want_result:
+    x: 5
+    y: 9.4
+- note: regocpp/expressions
+  modules:
+  - |
+    package expressions
+    
+    a := 5
+    b := -2
+    c := a + b
+    d := c * -2.5
+    e := (a + b) == 7
+    f := (c + d) != 7
+    g := 3 > -1
+    h := 2.5 >= a
+    i := b < a
+    j := 1.0 <= 2.1
+    k := a - (b + -c) * (-d / 2) + 3.7
+  query: data.expressions.k > 27 = x
+  want_result:
+    x: true
+- note: regocpp/bodies
+  modules:
+  - |
+    package bodies
+    
+    val := -1
+    a := 5
+    b := 3
+    c := a + b { a; b; a > b }
+    d := val {
+      input.a < a
+      val := a + data.bodies.b
+    } {
+      input.a >= data.bodies.a
+      val := a - b
+    }
+    
+    e := {"one": x, "two": y + z} {
+      x := a * b
+      c
+      y := d
+      z := x - y
+    }
+    
+    default f := false
+    
+    f {
+      b < a
+      not d
+    }
+    
+    f {
+      input.c > d
+    }
+    
+    g {
+      d == a + b
+      a < b
+    }
+  input:
+    a: 10
+    b: "20"
+    c: 30.0
+    d: true
+  query: data.bodies.e = x
+  want_result:
+    x:
+      one: 15
+      two: 15
+- note: regocpp/bodies_undef
+  modules:
+  - |
+    package bodies
+    
+    val := -1
+    a := 5
+    b := 3
+    c := a + b { a; b; a > b }
+    d := val {
+      input.a < a
+      val := a + data.bodies.b
+    } {
+      input.a >= data.bodies.a
+      val := a - b
+    }
+    
+    e := {"one": x, "two": y + z} {
+      x := a * b
+      c
+      y := d
+      z := x - y
+    }
+    
+    default f := false
+    
+    f {
+      b < a
+      not d
+    }
+    
+    f {
+      input.c > d
+    }
+    
+    g {
+      d == a + b
+      a < b
+    }
+  input:
+    a: 10
+    b: "20"
+    c: 30.0
+    d: true
+  query: data.bodies.g = x
+  want_defined: false
+- note: regocpp/sets
+  modules:
+  - |
+    package sets
+    
+    a := 1
+    b := 1
+    c := 2
+    d := 3
+    
+    e := "one"
+    f := "one"
+    g := "two"
+    h := "three" {
+      d > c
+      c > b
+      b >= a
+    }
+    
+    i := {"foo": 1}
+    j := {"foo": 1}
+    k := {"foo": 2, "bar": 1, "baz": 1}
+    l := {"foo": 3, "bar": 2, "baz": 1}
+    
+    int_set := {4, a, c, d, b}
+    string_set := {"four", g, f, e}
+    obj_set := {i, j, k, l}
+    array_set := {[a, b], [b, a], [a, c], [a, d], [b, c], [b, d], [c, d]}
+    
+    empty_set := set()
+    
+    m {
+      data.sets.int_set[1]
+      string_set["one"]
+      data.sets.obj_set[{"foo": 1}]
+      array_set[[1, 1]]
+    }
+  query: data.sets.string_set = x
+  want_result:
+    x:
+      four: null
+      one: null
+      two: null
+- note: regocpp/sets_membership
+  modules:
+  - |
+    package sets
+    
+    a := 1
+    b := 1
+    c := 2
+    d := 3
+    
+    e := "one"
+    f := "one"
+    g := "two"
+    h := "three" {
+      d > c
+      c > b
+      b >= a
+    }
+    
+    i := {"foo": 1}
+    j := {"foo": 1}
+    k := {"foo": 2, "bar": 1, "baz": 1}
+    l := {"foo": 3, "bar": 2, "baz": 1}
+    
+    int_set := {4, a, c, d, b}
+    string_set := {"four", g, f, e}
+    obj_set := {i, j, k, l}
+    array_set := {[a, b], [b, a], [a, c], [a, d], [b, c], [b, d], [c, d]}
+    
+    empty_set := set()
+    
+    m {
+      data.sets.int_set[1]
+      string_set["one"]
+      data.sets.obj_set[{"foo": 1}]
+      array_set[[1, 1]]
+    }
+  query: data.sets.m = x
+  want_result:
+    x: true
+- note: regocpp/var_keys
+  modules:
+  - |
+    package objects
+
+    rect := {`width`: 2, "height": 4}
+    cube := {"width": 3, `height`: 4, "depth": 5}
+    a := 42
+    b := false
+    c := null
+    d := {"a": a, "x": [b, c]}
+    index := 1
+    shapes := [rect, cube]
+    names := ["prod", `smoke1`, "dev"]
+    sites := [{"name": "prod"}, {"name": names[index]}, {"name": "dev"}]
+    e := {
+        a: "foo",
+        "three": c,
+        names[2]: b,
+        "four": d,
+    }
+    f := e["dev"]
+  query: data.objects.f = x
+  want_result:
+    x: false
+- note: regocpp/multi/0
+  input:
+    a: 10
+  modules:
+  - |
+    package multi
+
+    default a := 0
+
+    a := val {
+        input.a > 0
+        input.a < 10
+        input.a % 2 == 1
+        val := input.a * 10
+    } {
+        input.a > 0
+        input.a < 10
+        input.a % 2 == 0
+        val := input.a * 10 + 1
+    }
+    
+    a := input.a / 10 {
+        input.a >= 10
+    }
+  query: data.multi.a = x
+  want_result:
+    x: 1
+- note: regocpp/multi/1
+  input:
+    a: 4
+  modules:
+  - |
+    package multi
+
+    default a := 0
+
+    a := val {
+        input.a > 0
+        input.a < 10
+        input.a % 2 == 1
+        val := input.a * 10
+    } {
+        input.a > 0
+        input.a < 10
+        input.a % 2 == 0
+        val := input.a * 10 + 1
+    }
+    
+    a := input.a / 10 {
+        input.a >= 10
+    }
+  query: data.multi.a = x
+  want_result:
+    x: 41
+- note: regocpp/multi/2
+  input:
+    a: 7
+  modules:
+  - |
+    package multi
+
+    default a := 0
+
+    a := val {
+        input.a > 0
+        input.a < 10
+        input.a % 2 == 1
+        val := input.a * 10
+    } {
+        input.a > 0
+        input.a < 10
+        input.a % 2 == 0
+        val := input.a * 10 + 1
+    }
+    
+    a := input.a / 10 {
+        input.a >= 10
+    }
+  query: data.multi.a = x
+  want_result:
+    x: 70
+- note: regocpp/functions/sum
+  input:
+    key: sum
+  modules:
+  - |
+    package functions
+
+    f(a, b, c) := sum {
+      b != -1
+        g(input.key)
+        sum := a + b + c
+    }
+
+    f(a, -1, c) := sum {
+        g(input.key)
+        sum := a + c
+    }
+
+    f(a, b, c) := prod {
+        not g(input.key)
+        prod := a * b * c
+    }
+
+    g("sum") {
+        true
+    }
+
+    g("prod") {
+        false
+    }
+
+    h(a, b) := result {
+        result := {
+            "a": a,
+            "b": b
+        }
+    }
+  query: "[data.functions.f(1, 2, 3), data.functions.f(1, -1, 3)] = x"
+  want_result:
+    x:
+    - 6
+    - 4
+- note: regocpp/functions/prod
+  input:
+    key: prod
+  modules:
+  - |
+    package functions
+
+    f(a, b, c) := sum {
+      b != -1
+        g(input.key)
+        sum := a + b + c
+    }
+
+    f(a, -1, c) := sum {
+        g(input.key)
+        sum := a + c
+    }
+
+    f(a, b, c) := prod {
+        not g(input.key)
+        prod := a * b * c
+    }
+
+    g("sum") {
+        true
+    }
+
+    g("prod") {
+        false
+    }
+
+    h(a, b) := result {
+        result := {
+            "a": a,
+            "b": b
+        }
+    }
+  query: "[data.functions.f(1, 2, 3), data.functions.f(1, -1, 3)] = x"
+  want_result:
+    x:
+    - 6
+    - -3
+- note: regocpp/functions/ref
+  modules:
+  - |
+    package functions
+
+    f(a, b, c) := sum {
+      b != -1
+        g(input.key)
+        sum := a + b + c
+    }
+
+    f(a, -1, c) := sum {
+        g(input.key)
+        sum := a + c
+    }
+
+    f(a, b, c) := prod {
+        not g(input.key)
+        prod := a * b * c
+    }
+
+    g("sum") {
+        true
+    }
+
+    g("prod") {
+        false
+    }
+
+    h(a, b) := result {
+        result := {
+            "a": a,
+            "b": b
+        }
+    }
+  query: "[data.functions.h(1, 2).a, data.functions.h(3, 4).b] = x"
+  want_result:
+    x:
+    - 1
+    - 4
+- note: regocpp/unify
+  modules:
+  - |
+    package unify
+
+    tuples = [[1, 1], [1, 2], [2, 3], [3, 5], [5, 8]]
+    letters = {["a", 1], ["b", 2], ["c", 3], ["d", 4]}
+    mixed = [[0, 1, 2], [3, 4, 5, 6], {"foo": 7}]
+
+    a := {
+      "foo": 1,
+        "bar": 2,
+    }
+
+    b := x + y {
+        some x, y
+        [1, x] = tuples[i]
+        {"foo": y, "bar": 2} = a
+        i != 0
+    }
+
+    c {
+        i != 0
+        x = tuples[i]
+        i = x[1] - 1
+    }
+
+    d {
+      tuple := tuples[i]
+        tuple[0] >= 3
+        tuple[1] >= 5
+    }
+
+    e := key {
+        target = 2
+        some key
+      a[key] = target
+    }
+
+    f := x + y {
+        letters[["a", x]]
+        letters[["c", y]]
+    }
+
+    g := j {
+      x = mixed[i]
+        x[j] = 6
+    }
+
+    h := j {
+      x = mixed[i]
+        x[j] = 7
+    }
+
+    output := [a, b, c, d, e, f, g, h]
+  query: data.unify.output = x
+  want_result:
+    x:
+    - {"bar": 2, "foo": 1}
+    - 3
+    - true
+    - true
+    - bar
+    - 4
+    - 3
+    - foo
+- note: regocpp/compare
+  modules:
+  - |
+    package compare
+
+    a := {"one": 1, "two": 2}
+    b := {"two": 2, "three": 3}
+    c := "2"
+    d := "0001"
+    e := [0, 1, 2]
+    f := [0, 0, 0, 4]
+    g := {"b": "2", "a": "1"}
+    h := {"b": "1", "a": "2"}
+
+    output := [a < b, c < d, e < f, g < h]
+  query: data.compare.output = x
+  want_result:
+    x: [true, false, false, true]
+- note: regocpp/keywords
+  data:
+    sites:
+    - region: east
+      name: prod
+      servers:
+      - {"name": "web-0", "hostname": "hydrogen"}
+      - {"name": "web-1", "hostname": "helium"}
+      - {"name": "db-0", "hostname": "lithium"}
+    - region: west
+      name: smoke
+      servers:
+      - {"name": "web-1000", "hostname": "beryllium"}
+      - {"name": "web-1001", "hostname": "boron"}
+      - {"name": "db-1000", "hostname": "carbon"}
+    - region: west
+      name: dev
+      servers:
+      - {"name": "web-dev", "hostname": "nitrogen"}
+      - {"name": "db-dev", "hostname": "oxygen"}
+    apps:
+    - name: web
+      servers: ["web-0", "web-1", "web-1000", "web-1001", "web-dev"]
+    - name: mysql
+      servers: ["db-0", "db-1000"]
+    - name: mongodb
+      servers: ["db-dev"]
+    containers:
+    - image: redis
+      ipaddress: "10.0.0.1"
+      name: big_stallman
+    - image: nginx
+      ipaddress: "10.0.0.2"
+      name: cranky_euclid
+  input:
+    a: 10
+    b: "20"
+    c: 30.0
+    d: true
+  modules:
+  - |
+    package keywords
+
+    import future.keywords
+    import data.apps as apps
+    import data.sites as sites
+
+    letters := ["a", "b", "c", "d"]
+
+    a := 3
+    b := 5
+    c := 4 if input.a > a else := 5 if {input.a > a} else := -1 if input.a <= a
+
+    d := val if {
+        some i, x in letters
+        x == "b"
+        some j in [1, 2, "a"]
+        x == letters[j]
+        i == j
+        "d" in letters
+        val := "found"
+    }
+
+    e contains "0" if {
+        input.d
+        input.b == "20"
+    }
+
+    e contains "a" if {
+        input.a > a
+    }
+
+    e contains "1" if input.key == "sum"
+
+    e contains a + b if input.key == "sum"
+
+    e contains "2" if input.key == "prod"
+
+    e contains a * b if input.key == "prod"
+
+    f["a"] := a
+    f["b"] := b
+    f["c"] if c
+    f["sum"] := sum {
+        sum := a + b
+    }
+
+    app_to_hostnames_obj[app_name] := hostnames if {
+        app := apps[_]
+        app_name := app.name
+        hostnames := [hostname | name := app.servers[_]
+                                s := sites[_].servers[_]
+                                s.name == name
+                                hostname := s.hostname]
+    }
+
+    output := [c, d, e, f, app_to_hostnames_obj["mongodb"] == ["oxygen"]]
+  query: data.keywords.output = x
+  want_result:
+    x:
+    - 4
+    - "found"
+    - ["0", "a"]
+    - {"a": 3, "b": 5, "c": true, "sum": 8}
+    - true
+- note: regocpp/imports
+  modules:
+  - |
+    package scalars
+
+    greeting := "Hello"
+    max_height := 42
+    pi := 3.14159
+    allowed := true
+    location := null
+  - |
+    package scalars.more.stuff
+
+    import data.scalars.max_height
+    import data.scalars.more.stuff["and.more"].stuff.shared as shared_stuff
+    
+    shared := 7
+    a := max_height
+    b := a * shared + shared_stuff
+  - |
+    package scalars.more.stuff["and.more"].stuff
+
+    import data.scalars.max_height
+    import data.scalars["more"].stuff["shared"] as shared_stuff
+
+    shared := 11
+    a := max_height
+    b := a * shared + shared_stuff
+  - |
+    package scalars
+
+    output := [
+        data.scalars.more.stuff.b,
+        allowed,
+        data.scalars.more.stuff["and.more"].stuff.b
+    ]
+  query: data.scalars.output = x
+  want_result:
+    x:
+    - 305
+    - true
+    - 469
+- note: regocpp/with
+  modules:
+  - |
+    package bodies
+
+    val := -1
+    a := 5
+    b := 3
+    c := a + b { a; b; a > b }
+    d := val {
+      input.a < a
+      val := a + data.bodies.b
+    } {
+      input.a >= data.bodies.a
+      val := a - b
+    }
+
+    e := {"one": x, "two": y + z} {
+      x := a * b
+      c
+      y := d
+      z := x - y
+    }
+
+    default f := false
+
+    f {
+      b < a
+      not d
+    }
+
+    f {
+      input.c > d
+    }
+
+    g {
+      d == a + b
+      a < b
+    }
+  - |
+    package bodies
+
+    import future.keywords.if
+    import future.keywords.in
+
+    d_sum := result {
+        result := d with input.a as a - 1
+    }
+
+    d_diff := result {
+        result := d with input.a as a + 1
+    }
+
+    inner := [x, y] if {
+        x := input.foo
+        y := input.bar
+    }
+
+    middle := [a, b] if {
+        a := inner with input.foo as 100
+        b := input
+    }
+
+    outer := result if {
+        result := middle with input as {"foo": 200, "bar": 300}
+    }
+
+    mycount(x) := count(x)
+
+    mock_count(x) := 0 if "x" in x
+    mock_count(x) := count(x) if not "x" in x
+
+    use_count := result {
+        result := mycount([1, 2, 3]) with count as mock_count
+    }
+
+    use_mock := result {
+        result := mycount(["x", "y", "z"]) with count as mock_count
+    }
+
+    output := [d_sum, d_diff, outer, use_count, use_mock]
+  query: data.bodies.output = x
+  want_result:
+    x:
+    - 8
+    - 2
+    - [[100, 300], {"bar": 300, "foo": 200}]
+    - 3
+    - 0
+- note: regocpp/every
+  data:
+    sites:
+    - region: east
+      name: prod
+      servers:
+      - {"name": "web-0", "hostname": "hydrogen"}
+      - {"name": "web-1", "hostname": "helium"}
+      - {"name": "db-0", "hostname": "lithium"}
+    - region: west
+      name: smoke
+      servers:
+      - {"name": "web-1000", "hostname": "beryllium"}
+      - {"name": "web-1001", "hostname": "boron"}
+      - {"name": "db-1000", "hostname": "carbon"}
+    - region: west
+      name: dev
+      servers:
+      - {"name": "web-dev", "hostname": "nitrogen"}
+      - {"name": "db-dev", "hostname": "oxygen"}
+    apps:
+    - name: web
+      servers: ["web-0", "web-1", "web-1000", "web-1001", "web-dev"]
+    - name: mysql
+      servers: ["db-0", "db-1000"]
+    - name: mongodb
+      servers: ["db-dev"]
+    containers:
+    - image: redis
+      ipaddress: "10.0.0.1"
+      name: big_stallman
+    - image: nginx
+      ipaddress: "10.0.0.2"
+      name: cranky_euclid
+  modules:
+  - |
+    package every
+
+    import future.keywords.every
+    import future.keywords.if
+    import data.sites as sites
+
+    names_with_dev if {
+        some site in sites
+        site.name == "dev"
+
+        every server in site.servers {
+            endswith(server.name, "-dev")
+        }
+    }
+
+    array_domain if {
+        every i, x in [1, 2, 3] { x-i == 1 } # array domain
+    }
+
+    object_domain if {
+        every k, v in {"foo": "bar", "fox": "baz" } { # object domain
+            startswith(k, "f")
+            startswith(v, "b")
+        }
+    }
+
+    set_domain if {
+        every x in {1, 2, 3} { x != 4 } # set domain
+    }
+
+    xs := [2, 2, 4, 8]
+    larger_than_one(x) := x > 1
+
+    rule_every if {
+        every x in xs { larger_than_one(x) }
+    }
+
+    not_less_or_equal_one if not lte_one
+
+    lte_one if {
+        some x in xs
+        not larger_than_one(x)
+    }
+
+    output {
+      names_with_dev
+        array_domain
+        object_domain
+        set_domain
+        rule_every
+        not_less_or_equal_one
+    }
+  query: data.every.output = x
+  want_result:
+    x: true
+- note: regocpp/compr
+  data:
+    sites:
+    - region: east
+      name: prod
+      servers:
+      - {"name": "web-0", "hostname": "hydrogen"}
+      - {"name": "web-1", "hostname": "helium"}
+      - {"name": "db-0", "hostname": "lithium"}
+    - region: west
+      name: smoke
+      servers:
+      - {"name": "web-1000", "hostname": "beryllium"}
+      - {"name": "web-1001", "hostname": "boron"}
+      - {"name": "db-1000", "hostname": "carbon"}
+    - region: west
+      name: dev
+      servers:
+      - {"name": "web-dev", "hostname": "nitrogen"}
+      - {"name": "db-dev", "hostname": "oxygen"}
+    apps:
+    - name: web
+      servers: ["web-0", "web-1", "web-1000", "web-1001", "web-dev"]
+    - name: mysql
+      servers: ["db-0", "db-1000"]
+    - name: mongodb
+      servers: ["db-dev"]
+    containers:
+    - image: redis
+      ipaddress: "10.0.0.1"
+      name: big_stallman
+    - image: nginx
+      ipaddress: "10.0.0.2"
+      name: cranky_euclid
+  modules:
+  - |
+    package compr
+
+    import future.keywords.if
+    import data.apps as apps
+    import data.sites as sites
+
+    region := "west"
+    names := [name | sites[i].region == region; name := sites[i].name]
+
+    app_to_hostnames := {app.name: hostnames |
+        app := apps[_]
+        hostnames := [hostname |
+                        name := app.servers[_]
+                        s := sites[_].servers[_]
+                        s.name == name
+                        hostname := s.hostname]
+    }
+
+    a := [1, 2, 3, 4, 3, 4, 3, 4, 5]
+    b := {x | x = a[_]}
+
+    output {
+        names == ["smoke", "dev"]
+        app_to_hostnames == {
+          "mongodb": ["oxygen"],
+            "mysql": ["lithium","carbon"],
+            "web": ["hydrogen", "helium", "beryllium", "boron", "nitrogen"]
+        }
+        b == {1, 2, 3, 4, 5}
+    }
+  query: data.compr.output = x
+  want_result:
+    x: true
+- note: regocpp/setops
+  modules:
+  - |
+    package setops
+
+    import future.keywords.in
+
+    a := {1, 2, 3, 4}
+    b := {3, 4, 5}
+    c := {4, 5, 6}
+    d := {2, 3, 4}
+    e := a | b & c | d
+    f := {x | some x in b & c | d}
+    g := [a | b]
+    h := [a | b, 4]
+    i := { a: b | c }
+    j := { a: b, c: d | e}
+    k := a - d
+
+    output {
+      e == {1, 2, 3, 4, 5}
+        f == {2, 3, 4, 5}
+        g == [a]
+        h == [{1, 2, 3, 4, 5}, 4]
+        i == {
+          a: {3, 4, 5}
+        }
+        j == {
+          a: {3, 4, 5},
+            c: {1, 2, 3, 4, 5}
+        }
+        k == {1}
+        1 in a | b
+        3 in a & b | c & d
+    }
+  query: data.setops.output = x
+  want_result:
+    x: true
+- note: regocpp/enumerate
+  modules:
+  - |
+    package enumerate
+
+    import future.keywords.in
+
+    numbers := [1, 2, 3, 4, 5]
+    letters := ["1", "2", "3", "4", "5"]
+
+    a := item {
+        some n in numbers
+        some l in letters
+        to_number(l) == n
+        item = [n, l]
+        item[0] == 3
+    }
+  query: data.enumerate.a = x
+  want_result:
+    x:
+      - 3
+      - "3"
diff --git a/tests/test_case.cc b/tests/test_case.cc
new file mode 100644
index 00000000..cfc5c1f5
--- /dev/null
+++ b/tests/test_case.cc
@@ -0,0 +1,549 @@
+#include "test_case.h"
+
+#include "rego_test.h"
+
+namespace
+{
+  using namespace rego_test;
+
+  void write_ast(
+    const std::filesystem::path& debug_path,
+    std::size_t index,
+    const std::string& pass,
+    const Node& ast)
+  {
+    if (debug_path.empty())
+    {
+      return;
+    }
+
+    if (!std::filesystem::is_directory(debug_path))
+    {
+      std::filesystem::create_directory(debug_path);
+    }
+
+    std::filesystem::path output;
+    if (index < 10)
+    {
+      output =
+        debug_path / ("0" + std::to_string(index) + "_" + pass + ".trieste");
+    }
+    else
+    {
+      output = debug_path / (std::to_string(index) + "_" + pass + ".trieste");
+    }
+    std::ofstream f(output, std::ios::binary | std::ios::out);
+
+    if (f)
+    {
+      // Write the AST to the output file.
+      f << "rego_test_cases" << std::endl << pass << std::endl << ast;
+      return;
+    }
+
+    std::cerr << "Could not open " << output << " for writing." << std::endl;
+  }
+
+  bool has_error(const Node& node)
+  {
+    if (node->type() == Error)
+    {
+      return true;
+    }
+
+    for (auto& child : *node)
+    {
+      if (has_error(child))
+      {
+        return true;
+      }
+    }
+
+    return false;
+  }
+
+  std::optional<Node> maybe_get_file(Node mapping, const std::string& name)
+  {
+    Location loc(name);
+    Nodes defs = mapping->lookdown(loc);
+    if (defs.size() > 1)
+    {
+      throw std::runtime_error("Expected exactly one " + name + " block");
+    }
+
+    if (defs.size() == 1)
+    {
+      Node val = defs[0]->back();
+      assert(val->type() == File);
+      return val;
+    }
+
+    return std::nullopt;
+  }
+
+  std::optional<std::string> maybe_get_string(
+    Node mapping, const std::string& name)
+  {
+    Location loc(name);
+    Nodes defs = mapping->lookdown(loc);
+    if (defs.size() > 1)
+    {
+      throw std::runtime_error("Expected exactly one " + name + " block");
+    }
+
+    if (defs.size() == 1)
+    {
+      Node val = defs[0]->back();
+      assert(val->type() == rego_test::Scalar);
+      val = val->front();
+      assert(val->type() == rego_test::String);
+      return std::string(val->location().view());
+    }
+
+    return std::nullopt;
+  }
+
+  std::string get_string(Node mapping, const std::string& name)
+  {
+    auto maybe_string = maybe_get_string(mapping, name);
+    if (maybe_string.has_value())
+    {
+      return *maybe_string;
+    }
+
+    return "";
+  }
+
+  Node get_node(Node mapping, const std::string& name)
+  {
+    Location loc(name);
+    Nodes defs = mapping->lookdown(loc);
+    if (defs.size() > 1)
+    {
+      throw std::runtime_error("Expected exactly one " + name + " block");
+    }
+
+    if (defs.size() == 1)
+    {
+      return defs[0]->back();
+    }
+
+    return {};
+  }
+
+  bool get_bool(Node mapping, const std::string& name)
+  {
+    Location loc(name);
+    Nodes defs = mapping->lookdown(loc);
+    if (defs.size() > 1)
+    {
+      throw std::runtime_error("Expected exactly one " + name + " block");
+    }
+
+    if (defs.size() == 1)
+    {
+      Node val = defs[0]->back();
+      assert(val->type() == rego_test::Scalar);
+      val = val->front();
+      assert(val->type() == rego_test::True || val->type() == rego_test::False);
+      return val->type() == rego_test::True;
+    }
+
+    return false;
+  }
+
+  std::vector<std::string> get_modules(Node mapping)
+  {
+    Location loc("modules");
+    Nodes defs = mapping->lookdown(loc);
+    if (defs.size() > 1)
+    {
+      throw std::runtime_error("Expected exactly one modules block");
+    }
+
+    std::vector<std::string> modules;
+    if (defs.size() == 1)
+    {
+      Node module_nodes = defs[0]->back();
+      assert(module_nodes->type() == rego_test::Sequence);
+      for (auto& entry : *module_nodes)
+      {
+        Node scalar = entry->front();
+        assert(scalar->type() == rego_test::Scalar);
+        Node module = scalar->front();
+        assert(module->type() == rego_test::String);
+        modules.push_back(std::string(module->location().view()));
+      }
+    }
+
+    return modules;
+  }
+
+  using BindingMap = std::map<std::string, std::string>;
+
+  BindingMap to_binding_map(const Node& node)
+  {
+    BindingMap bindings;
+    for (auto& binding : *node)
+    {
+      std::string key = std::string((binding / rego::Var)->location().view());
+      std::string value = rego::to_json((binding / rego::Term));
+      bindings[key] = value;
+    }
+
+    return bindings;
+  }
+
+  void diff(std::string actual, std::string wanted, std::ostream& os)
+  {
+    std::set<std::size_t> errors;
+    std::size_t index = 0;
+    std::size_t length = std::min(actual.size(), wanted.size());
+    for (; index < length; ++index)
+    {
+      if (actual[index] != wanted[index])
+      {
+        errors.insert(index);
+      }
+    }
+
+    length = std::max(actual.size(), wanted.size());
+    for (; index < length; ++index)
+    {
+      errors.insert(index);
+    }
+
+    os << "wanted: " << wanted << std::endl;
+    os << "actual: " << actual << std::endl;
+    os << "        ";
+    for (std::size_t i = 0; i < length; ++i)
+    {
+      if (errors.contains(i))
+      {
+        os << "^";
+      }
+      else
+      {
+        os << " ";
+      }
+    }
+  }
+
+  bool compare(Node actual, Node wanted, std::ostream& os)
+  {
+    BindingMap actual_bindings = to_binding_map(actual);
+    BindingMap wanted_bindings = to_binding_map(wanted);
+    for (auto& [key, value] : wanted_bindings)
+    {
+      if (!actual_bindings.contains(key))
+      {
+        os << "Missing binding for " << key << std::endl;
+        return false;
+      }
+
+      if (actual_bindings[key] != value)
+      {
+        diff(key + "= " + actual_bindings[key], key + "= " + value, os);
+        return false;
+      }
+    }
+
+    return true;
+  }
+}
+
+namespace rego_test
+{
+  std::optional<std::vector<TestCase>> TestCase::load(
+    const std::filesystem::path& path, const std::filesystem::path& debug_path)
+  {
+    auto ast = parser().parse(path);
+    bool ok = wf_parser.build_st(ast, std::cerr);
+    ok = wf_parser.check(ast, std::cerr) && ok;
+
+    write_ast(debug_path, 0, "parse", ast);
+    if (!ok)
+    {
+      std::ostringstream buf;
+      ast->errors(buf);
+      throw std::runtime_error(buf.str());
+    }
+
+    auto passes = rego_test::passes();
+
+    for (std::size_t i = 0; i < passes.size(); ++i)
+    {
+      auto& [pass_name, pass, wf] = passes[i];
+      wf::push_back(wf);
+      auto [new_ast, count, changes] = pass->run(ast);
+      wf::pop_front();
+      ast = new_ast;
+
+      ok = wf->build_st(ast, std::cout);
+      write_ast(debug_path, i + 1, pass_name, ast);
+
+      ok = wf->check(ast, std::cout) && ok;
+
+      if (has_error(ast))
+      {
+        ok = false;
+      }
+
+      if (!ok)
+      {
+        std::ostringstream buf;
+        buf << "Failed at pass " << pass_name << std::endl;
+        ast->errors(buf);
+        std::cerr << buf.str() << std::endl;
+        return std::nullopt;
+      }
+    }
+
+    std::vector<TestCase> test_cases;
+    Node case_seq = ast->front()->front()->back();
+    for (Node entry : *case_seq)
+    {
+      Node test_case = entry->front();
+      test_cases.push_back(
+        TestCase::create_from_node(test_case).filename(path));
+    }
+
+    return test_cases;
+  }
+
+  TestCase TestCase::create_from_node(const Node& test_case_map)
+  {
+    TestCase test_case;
+    auto note = maybe_get_string(test_case_map, "note");
+    if (note.has_value())
+    {
+      test_case = test_case.note(*note);
+    }
+    else
+    {
+      throw std::runtime_error("Note is required");
+    }
+
+    auto query = maybe_get_string(test_case_map, "query");
+    if (query.has_value())
+    {
+      test_case = test_case.query(*query);
+    }
+    else
+    {
+      throw std::runtime_error("Query is required");
+    }
+
+    auto data = maybe_get_file(test_case_map, "data");
+    if (data.has_value())
+    {
+      test_case = test_case.data(*data);
+    }
+
+    auto input = maybe_get_file(test_case_map, "input");
+    if (input.has_value())
+    {
+      test_case = test_case.input(*input);
+    }
+
+    return test_case.modules(get_modules(test_case_map))
+      .input_term(get_string(test_case_map, "input_term"))
+      .want_defined(get_bool(test_case_map, "want_defined"))
+      .want_result(get_node(test_case_map, "want_result"))
+      .want_error_code(get_string(test_case_map, "want_error_code"))
+      .want_error(get_string(test_case_map, "want_error"))
+      .sort_bindings(get_bool(test_case_map, "sort_bindings"))
+      .strict_error(get_bool(test_case_map, "strict_error"));
+  }
+
+  Result TestCase::run(const std::filesystem::path& executable_path) const
+  {
+    rego::Interpreter interpreter;
+    interpreter.executable(executable_path);
+    for (std::size_t i = 0; i < m_modules.size(); ++i)
+    {
+      std::string name = "module" + std::to_string(i);
+      interpreter.add_module(name, m_modules[i]);
+    }
+    interpreter.add_data(m_data);
+    if (m_input_term.size() > 0)
+    {
+      interpreter.add_input_json(m_input_term);
+    }
+    else
+    {
+      interpreter.add_input(m_input);
+    }
+
+    bool pass = true;
+    std::ostringstream error;
+    Node actual = interpreter.raw_query(m_query);
+    if (m_want_error.length() > 0)
+    {
+      pass = false;
+      error << "Unsupported test option";
+    }
+    else if (m_want_error_code.length() > 0)
+    {
+      pass = false;
+      error << "Unsupported test option";
+    }
+    else if (m_want_result)
+    {
+      pass = compare(actual, m_want_result, error);
+    }
+    else if (m_want_defined)
+    {
+      pass = actual->front()->type() != Undefined;
+    }
+    else
+    {
+      pass = actual->front()->type() == Undefined;
+    }
+
+    return {pass, error.str()};
+  }
+
+  const std::filesystem::path& TestCase::filename() const
+  {
+    return m_filename;
+  }
+
+  TestCase& TestCase::filename(const std::filesystem::path& filename)
+  {
+    m_filename = filename;
+    return *this;
+  }
+
+  const std::string& TestCase::note() const
+  {
+    return m_note;
+  }
+
+  TestCase& TestCase::note(const std::string& note)
+  {
+    m_note = note;
+    return *this;
+  }
+
+  const std::string& TestCase::query() const
+  {
+    return m_query;
+  }
+
+  TestCase& TestCase::query(const std::string& query)
+  {
+    m_query = query;
+    return *this;
+  }
+
+  const std::vector<std::string>& TestCase::modules() const
+  {
+    return m_modules;
+  }
+
+  TestCase& TestCase::modules(const std::vector<std::string>& modules)
+  {
+    m_modules = modules;
+    return *this;
+  }
+
+  const Node& TestCase::data() const
+  {
+    return m_data;
+  }
+
+  TestCase& TestCase::data(const Node& data)
+  {
+    m_data = data;
+    return *this;
+  }
+
+  const Node& TestCase::input() const
+  {
+    return m_input;
+  }
+
+  TestCase& TestCase::input(const Node& input)
+  {
+    m_input = input;
+    return *this;
+  }
+
+  const std::string& TestCase::input_term() const
+  {
+    return m_input_term;
+  }
+
+  TestCase& TestCase::input_term(const std::string& input_term)
+  {
+    m_input_term = input_term;
+    return *this;
+  }
+
+  bool TestCase::want_defined() const
+  {
+    return m_want_defined;
+  }
+
+  TestCase& TestCase::want_defined(bool want_defined)
+  {
+    m_want_defined = want_defined;
+    return *this;
+  }
+
+  const Node& TestCase::want_result() const
+  {
+    return m_want_result;
+  }
+
+  TestCase& TestCase::want_result(const Node& want_result)
+  {
+    m_want_result = want_result;
+    return *this;
+  }
+
+  const std::string& TestCase::want_error_code() const
+  {
+    return m_want_error_code;
+  }
+
+  TestCase& TestCase::want_error_code(const std::string& want_error_code)
+  {
+    m_want_error_code = want_error_code;
+    return *this;
+  }
+
+  const std::string& TestCase::want_error() const
+  {
+    return m_want_error;
+  }
+
+  TestCase& TestCase::want_error(const std::string& want_error)
+  {
+    m_want_error = want_error;
+    return *this;
+  }
+
+  bool TestCase::sort_bindings() const
+  {
+    return m_sort_bindings;
+  }
+
+  TestCase& TestCase::sort_bindings(bool sort_bindings)
+  {
+    m_sort_bindings = sort_bindings;
+    return *this;
+  }
+
+  bool TestCase::strict_error() const
+  {
+    return m_strict_error;
+  }
+
+  TestCase& TestCase::strict_error(bool strict_error)
+  {
+    m_strict_error = strict_error;
+    return *this;
+  }
+}
\ No newline at end of file
diff --git a/tests/test_case.h b/tests/test_case.h
new file mode 100644
index 00000000..4df3d2ed
--- /dev/null
+++ b/tests/test_case.h
@@ -0,0 +1,97 @@
+#pragma once
+
+#include "rego.h"
+
+#include <string>
+#include <vector>
+
+namespace rego_test
+{
+  using namespace rego;
+
+  struct Result
+  {
+    bool passed;
+    std::string error;
+  };
+
+  class TestCase
+  {
+  public:
+    /** Loads a test case from a YAML file */
+    static std::optional<std::vector<TestCase>> load(
+      const std::filesystem::path& path,
+      const std::filesystem::path& debug_path = "");
+
+    Result run(const std::filesystem::path& executable_path) const;
+
+    /** name of file that case was loaded from */
+    const std::filesystem::path& filename() const;
+    TestCase& filename(const std::filesystem::path& filename);
+
+    /** globally unique identifier for this test case */
+    const std::string& note() const;
+    TestCase& note(const std::string& note);
+
+    /** policy query to execute */
+    const std::string& query() const;
+    TestCase& query(const std::string& query);
+
+    /** policies to test against */
+    const std::vector<std::string>& modules() const;
+    TestCase& modules(const std::vector<std::string>& modules);
+
+    /** data to tests against */
+    const Node& data() const;
+    TestCase& data(const Node& data);
+
+    /** parsed input data to use */
+    const Node& input() const;
+    TestCase& input(const Node& input);
+
+    /** raw input data (serialized to a string, overrides input) */
+    const std::string& input_term() const;
+    TestCase& input_term(const std::string& input_term);
+
+    /** expect query result to be defined (or not) */
+    bool want_defined() const;
+    TestCase& want_defined(bool want_defined);
+
+    /** expect query result (overrides defined) */
+    const Node& want_result() const;
+    TestCase& want_result(const Node& want_result);
+
+    /** expect query error code (overrides result) */
+    const std::string& want_error_code() const;
+    TestCase& want_error_code(const std::string& want_error_code);
+
+    /** expect query error message (overrides error code) */
+    const std::string& want_error() const;
+    TestCase& want_error(const std::string& want_error);
+
+    /** indicates that binding values should be treated as sets */
+    bool sort_bindings() const;
+    TestCase& sort_bindings(bool sort_bindings);
+
+    /** indicates that the error message depends on strict builtin error mode */
+    bool strict_error() const;
+    TestCase& strict_error(bool strict_error);
+
+  private:
+    static TestCase create_from_node(const Node& test_case_map);
+    std::filesystem::path m_filename;
+    std::string m_note;
+    std::string m_query;
+    std::vector<std::string> m_modules;
+    Node m_data;
+    Node m_input;
+    std::string m_input_term;
+    bool m_want_defined;
+    Node m_want_result;
+    std::string m_want_error_code;
+    std::string m_want_error;
+    bool m_sort_bindings;
+    bool m_strict_error;
+  };
+
+}