Fixing an issue with implicit enum statement capture (#103)

- Fixes an issue with implicit enum statement capture
- Fixes an issue when rule definitions are on a new line and are missing an assign operator

Signed-off-by: Matthew Johnson <matjoh@microsoft.com>

---------

Signed-off-by: Matthew Johnson <matjoh@microsoft.com>

include/rego/rego.hh

@@ -708,7 +708,7 @@ namespace rego
   inline const auto wf_pass_init =
     wf_pass_simple_refs
     | (UnifyBody <<= (Local | Literal | LiteralWith | LiteralEnum | LiteralNot | LiteralInit)++[1])
-    | (LiteralInit <<= VarSeq * VarSeq * AssignInfix)
+    | (LiteralInit <<= (Lhs >>= VarSeq) * (Rhs >>= VarSeq) * AssignInfix)
     ;
   // clang-format on
 

src/parse.cc

@@ -1,5 +1,4 @@
 #include "internal.hh"
-#include "rego.hh"
 
 namespace rego
 {

src/passes/enumerate.cc

@@ -65,6 +65,64 @@ namespace
     auto it = unifybody->find(local) + 1;
     return find_enum(unifybody, it);
   }
+
+  void vars_from(Node node, std::set<Location>& vars)
+  {
+    if (node->type() == Var)
+    {
+      vars.insert(node->location());
+    }
+
+    for (Node child : *node)
+    {
+      vars_from(child, vars);
+    }
+  }
+
+  // Determines which statements following an implicit enum statement are needed
+  // to instantiate the item sequence for that enum and adds them to the outside
+  // list. Otherwise they should be captured by the enum and are placed in the
+  // inside list.
+  void capture_statements(
+    const NodeRange tail,
+    Node itemseq,
+    std::vector<Node>& outside,
+    std::vector<Node>& inside)
+  {
+    std::set<Location> vars;
+    vars_from(itemseq, vars);
+    for (auto it = tail.first; it != tail.second; ++it)
+    {
+      Node stmt = *it;
+      if (stmt->type() == LiteralInit)
+      {
+        std::set<Location> init_vars;
+        vars_from(stmt / Lhs, init_vars);
+        vars_from(stmt / Rhs, init_vars);
+        std::set<Location> intersection;
+        std::set_intersection(
+          vars.begin(),
+          vars.end(),
+          init_vars.begin(),
+          init_vars.end(),
+          std::inserter(intersection, intersection.begin()));
+        if (intersection.empty())
+        {
+          inside.push_back(stmt);
+        }
+        else
+        {
+          // the item sequence depends on this init, so it must
+          // be outside the enum.
+          outside.push_back(stmt);
+        }
+      }
+      else
+      {
+        inside.push_back(stmt);
+      }
+    }
+  }
 }
 
 namespace rego
@@ -170,10 +228,14 @@ namespace rego
               return err(idx, "Invalid index for enumeration");
             }
 
+            std::vector<Node> outside;
+            std::vector<Node> inside;
+            capture_statements(_[Tail], _(ItemSeq), outside, inside);
+
             auto temp = _.fresh({"enum"});
             auto item = _.fresh({"item"});
             return Seq
-              << (Local << (Var ^ item) << Undefined)
+              << outside << (Local << (Var ^ item) << Undefined)
               << (LiteralEnum
                   << (Var ^ item) << _(ItemSeq)
                   << (UnifyBody
@@ -197,7 +259,7 @@ namespace rego
                                           << (Var ^ item)
                                           << (RefArgBrack
                                               << (Scalar << (Int ^ "1"))))))))
-                      << _[Tail]));
+                      << inside));
           },
 
         In(UnifyBody) *
@@ -240,10 +302,14 @@ namespace rego
               idx = Term << idx;
             }
 
+            std::vector<Node> outside;
+            std::vector<Node> inside;
+            capture_statements(_[Tail], _(ItemSeq), outside, inside);
+
             auto temp = _.fresh({"enum"});
             auto item = _.fresh({"item"});
             return Seq
-              << (Local << (Var ^ item) << Undefined)
+              << outside << (Local << (Var ^ item) << Undefined)
               << (LiteralEnum
                   << (Var ^ item) << _(ItemSeq)
                   << (UnifyBody
@@ -268,7 +334,7 @@ namespace rego
                                               << (RefArgBrack
                                                   << (Scalar
                                                       << (Int ^ "1")))))))))
-                      << _[Tail]));
+                      << inside));
           },
       }};
   }

src/passes/init.cc

@@ -3,7 +3,6 @@
 #include <algorithm>
 #include <cstddef>
 #include <deque>
-#include <stdexcept>
 
 namespace
 {

src/passes/rules.cc

@@ -317,6 +317,13 @@ namespace rego
             return Seq << _[RuleRef];
           },
 
+        In(Policy) * (T(Group) << T(Var)[Var]) *
+            (T(Group) << T(UnifyBody)[UnifyBody]) >>
+          [](Match& _) {
+            ACTION();
+            return Group << _(Var) << _(UnifyBody);
+          },
+
         // errors
 
         In(Policy) * T(Group)[Group] >>

tests/regocpp.yaml

@@ -1122,3 +1122,53 @@ cases:
       - 1
       - 2
       - 3
+- note: regocpp/bug101
+  modules:
+  - |
+    package test
+
+    x = y {
+      y = [z |
+        z = a[_]
+        a = [1, 2, 3]
+      ]
+    }
+  query: data.test.x = x
+  want_result:
+    - x:
+      - 1
+      - 2
+      - 3
+- note: regocpp/non-monotone
+  modules:
+  - |
+    package test
+
+    ips_by_port := {
+        80: ["1.1.1.1", "1.1.1.2"],
+        443: ["2.2.2.1"],
+    }
+
+    default foo1 := false
+    default foo2 := false
+
+    foo1
+        {
+        port = 77;
+        a = [port | ips_by_port[port][_] == "2.2.2.1" ];
+        b = [port | ips_by_port[port][_] == "1.1.1.1" ];
+        a = b
+        }
+
+    foo2
+        {
+        a = [port | ips_by_port[port][_] == "2.2.2.1" ];
+        b = [port | ips_by_port[port][_] == "1.1.1.1" ];
+        a = b
+        }
+
+  query: [data.test.foo1, data.test.foo2] = x
+  want_result:
+    - x:
+      - true
+      - false