Repository: microsoft/regorus (language: python)

Issue:
multiple underscore indexes issue
test case:

```
x1[y] {
  y = sites[_].servers[_].hostname
}
```

data from: https://www.openpolicyagent.org/docs/latest/policy-language/#example-data

error:
```
   |
57 |           y = sites[_].servers[_].hostname
   |                               ^
error: item cannot be indexed
```

I suspect this part of the code: 
https://github.com/microsoft/regorus/blob/main/src/interpreter.rs#L99

If I add variable with _ key / removed the condition on line 99, it's now working only for one underscore:
```
x1[y] {
  y = sites[i].servers[_].hostname
}
```

The OPA documentation mentions that:
> Under the hood, OPA translates the _ character to a unique variable name that does not conflict with variables and rules that are in scope.

I believe this is how it should be done, any hint to do that?



Discussion / hints:

Fix patch:
diff --git a/src/interpreter.rs b/src/interpreter.rs
index 3a95617f..64497c6a 100644
--- a/src/interpreter.rs
+++ b/src/interpreter.rs
@@ -122,6 +122,9 @@ impl<'source> Interpreter<'source> {
     }
 
     fn current_scope(&mut self) -> Result<&Scope> {
+        if self.scopes.is_empty() {
+            println!("here");
+        }
         self.scopes
             .last()
             .ok_or_else(|| anyhow!("internal error: no active scope"))
@@ -170,6 +173,10 @@ impl<'source> Interpreter<'source> {
         // Collect a chaing of '.field' or '["field"]'
         let mut path = vec![];
         loop {
+            if let Some(v) = self.loop_var_values.get(expr) {
+                path.reverse();
+                return Ok(Self::get_value_chained(v.clone(), &path[..]));
+            }
             match expr {
                 // Stop path collection upon encountering the leading variable.
                 Expr::Var(v) => {
@@ -388,7 +395,7 @@ impl<'source> Interpreter<'source> {
 
                         match (&lhs_var, &rhs_var) {
                             (Value::Undefined, Value::Undefined) => {
-                                bail!(lhs.span().error("both operators are unsafe"))
+                                bail!(lhs.span().error("both operands are unsafe"))
                             }
                             (Value::Undefined, _) => (lhs_name, rhs_var),
                             (_, Value::Undefined) => (rhs_name, lhs_var),
@@ -854,17 +861,30 @@ impl<'source> Interpreter<'source> {
             let loop_expr = &loops[0];
             let mut result = false;
 
+            let loop_expr_value = self.eval_expr(loop_expr.value)?;
+
+            // If the loop's index variable has already been assigned a value
+            // (this can happen if the same index is used for two different collections),
+            // then evaluate statements only if the index applies to this collection.
+            if let Some(idx) = self.lookup_local_var(loop_expr.index) {
+                if loop_expr_value[&idx] != Value::Undefined {
+                    result = self.eval_stmts_in_loop(stmts, &loops[1..])? || result;
+                }
+                return Ok(result);
+            }
+
             // Save the current scope and restore it after evaluating the statements so
             // that the effects of the current loop iteration are cleared.
             let scope_saved = self.current_scope()?.clone();
 
-            match self.eval_expr(loop_expr.value)? {
+            match loop_expr_value {
                 Value::Array(items) => {
                     for (idx, v) in items.iter().enumerate() {
                         self.loop_var_values.insert(loop_expr.expr, v.clone());
                         self.add_variable(loop_expr.index, Value::from_float(idx as Float))?;
 
                         result = self.eval_stmts_in_loop(stmts, &loops[1..])? || result;
+                        self.loop_var_values.remove(loop_expr.expr);
                         *self.current_scope_mut()? = scope_saved.clone();
                     }
                 }
@@ -874,6 +894,7 @@ impl<'source> Interpreter<'source> {
                         // For sets, index is also the value.
                         self.add_variable(loop_expr.index, v.clone())?;
                         result = self.eval_stmts_in_loop(stmts, &loops[1..])? || result;
+                        self.loop_var_values.remove(loop_expr.expr);
                         *self.current_scope_mut()? = scope_saved.clone();
                     }
                 }
@@ -883,6 +904,7 @@ impl<'source> Interpreter<'source> {
                         // For objects, index is key.
                         self.add_variable(loop_expr.index, k.clone())?;
                         result = self.eval_stmts_in_loop(stmts, &loops[1..])? || result;
+                        self.loop_var_values.remove(loop_expr.expr);
                         *self.current_scope_mut()? = scope_saved.clone();
                     }
                 }
@@ -930,7 +952,10 @@ impl<'source> Interpreter<'source> {
                             _ => map.insert(key, value),
                         };
                     } else {
-                        ctx.value = Value::Undefined;
+                        match &ctx.value {
+                            Value::Object(_) => (),
+                            _ => ctx.value = Value::Undefined,
+                        }
                     };
                 }
                 (None, Some(oe)) => {
@@ -947,7 +972,10 @@ impl<'source> Interpreter<'source> {
                             _ => bail!("internal error: invalid context value"),
                         }
                     } else {
-                        ctx.value = Value::Undefined;
+                        match &ctx.value {
+                            Value::Set(_) => (),
+                            _ => ctx.value = Value::Undefined,
+                        }
                     }
                 }
                 // No output expression.
@@ -1512,10 +1540,7 @@ impl<'source> Interpreter<'source> {
             // TODO: Handle undefined variables
             Expr::Var(_) => self.eval_chained_ref_dot_or_brack(expr),
             Expr::RefDot { .. } => self.eval_chained_ref_dot_or_brack(expr),
-            Expr::RefBrack { .. } => match self.loop_var_values.get(expr) {
-                Some(v) => Ok(v.clone()),
-                _ => self.eval_chained_ref_dot_or_brack(expr),
-            },
+            Expr::RefBrack { .. } => self.eval_chained_ref_dot_or_brack(expr),
 
             // Expressions with operators
             Expr::ArithExpr { op, lhs, rhs, .. } => self.eval_arith_expr(op, lhs, rhs),
@@ -1603,6 +1628,7 @@ impl<'source> Interpreter<'source> {
         span: &'source Span<'source>,
         bodies: &'source Vec<RuleBody<'source>>,
     ) -> Result<Value> {
+        let n_scopes = self.scopes.len();
         let result = if bodies.is_empty() {
             self.contexts.push(ctx.clone());
             self.eval_output_expr()
@@ -1617,7 +1643,6 @@ impl<'source> Interpreter<'source> {
                 }
 
                 // TODO: Manage other scoped data.
-                self.scopes.pop();
                 if bodies.len() > 1 {
                     unimplemented!("else bodies");
                 }
@@ -1635,8 +1660,7 @@ impl<'source> Interpreter<'source> {
             Err(e) => return Err(e),
         };
 
-        // Drop local variables and leave the local scope
-        self.scopes.pop();
+        assert_eq!(self.scopes.len(), n_scopes);
 
         Ok(match result {
             true => match &ctx.value {
@@ -1651,7 +1675,7 @@ impl<'source> Interpreter<'source> {
                     ))
                 }
                 Value::Set(_) => ctx.value,
-                _ => unimplemented!("todo fix this"),
+                _ => unimplemented!("todo fix this: ctx.value = {:?}", ctx.value),
             },
             false => Value::Undefined,
         })
diff --git a/src/scheduler.rs b/src/scheduler.rs
index 0db11019..1415b194 100644
--- a/src/scheduler.rs
+++ b/src/scheduler.rs
@@ -661,7 +661,7 @@ impl<'a> Analyzer<'a> {
                     )?;
                     self.locals.get(query)
                 }
-                _ => return Ok(()),
+                _ => break,
             };
 
             // Record vars used by the comprehension scope.
diff --git a/src/value.rs b/src/value.rs
index bfb21456..ec0678ce 100644
--- a/src/value.rs
+++ b/src/value.rs
@@ -344,6 +344,10 @@ impl ops::Index<&Value> for Value {
                 Some(v) => v,
                 _ => &Value::Undefined,
             },
+            (Value::Set(s), _) => match s.get(key) {
+                Some(v) => v,
+                _ => &Value::Undefined,
+            },
             (Value::Array(a), Value::Number(n)) => {
                 let index = n.0 .0 as usize;
                 if index < a.len() {


Test patch:
diff --git a/tests/interpreter/cases/loop/chained.yaml b/tests/interpreter/cases/loop/chained.yaml
new file mode 100644
index 00000000..107444b2
--- /dev/null
+++ b/tests/interpreter/cases/loop/chained.yaml
@@ -0,0 +1,152 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+cases:
+  - note: chained
+    data: {}
+    modules:
+      - |
+        package test
+
+        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"
+            }
+        ]
+        
+        x1[y] {
+          y = sites[_].servers[_].hostname
+        }
+
+        x5[y] {
+          y = sites[i].servers[i].hostname
+        }
+
+        obj = {
+         "a" : {
+           "a" : "b",
+           "b" : "c"
+           },
+          "c" : {
+            "c" : "d"
+          }
+        }
+        
+        x6[y] {
+         y = obj[i][i]
+        }
+
+        # Another definition for x6
+        x6[y] {
+         y = obj[i][obj[i][i]]
+        }
+        
+        x7[y] {
+         y = obj[i][obj[i][i]]
+        }
+        
+        results = {
+          "x1" : x1,
+          "x2" : x1 == { y | y = sites[i].servers[_].hostname },
+          "x3" : x1 == { y | y = sites[_].servers[i].hostname },
+          "x4" : x1 == { y | y = sites[i].servers[j].hostname },
+          "x6" : x6,
+          "x7" : x7,
+        }
+
+
+    query: data.test.results
+    want_result:
+      x1:
+        set!: [
+        "beryllium",
+        "boron",
+        "carbon",
+        "helium",
+        "hydrogen",
+        "lithium",
+        "nitrogen",
+        "oxygen",
+        ]
+      x2: true
+      x3: true
+      x4: true
+      x6:
+        set!: ["b", "c", "d"]
+      x7:
+        set!: ["c"]
+