From 7e3fc08a14b85fe40a0d28f59932c1ac92b8f937 Mon Sep 17 00:00:00 2001
From: Anand Krishnamoorthi <35780660+anakrish@users.noreply.github.com>
Date: Fri, 15 Mar 2024 07:13:54 -0700
Subject: [PATCH] Handle non simple refs in chained expressions (#182)

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>
---
 src/interpreter.rs                      | 14 ++++++--------
 src/utils.rs                            |  3 ++-
 tests/interpreter/cases/refr/tests.yaml | 14 ++++++++++++++
 3 files changed, 22 insertions(+), 9 deletions(-)
 create mode 100644 tests/interpreter/cases/refr/tests.yaml

diff --git a/src/interpreter.rs b/src/interpreter.rs
index 677fbf8b..3e9434ee 100644
--- a/src/interpreter.rs
+++ b/src/interpreter.rs
@@ -2926,25 +2926,23 @@ impl Interpreter {
     pub fn get_path_string(refr: &Expr, document: Option<&str>) -> Result<String> {
         let mut comps = vec![];
         let mut expr = Some(refr);
-        while expr.is_some() {
-            match expr {
-                Some(Expr::RefDot { refr, field, .. }) => {
+        while let Some(e) = expr {
+            match e {
+                Expr::RefDot { refr, field, .. } => {
                     comps.push(field.text());
                     expr = Some(refr);
                 }
-                Some(Expr::RefBrack { refr, index, .. })
-                    if matches!(index.as_ref(), Expr::String(_)) =>
-                {
+                Expr::RefBrack { refr, index, .. } if matches!(index.as_ref(), Expr::String(_)) => {
                     if let Expr::String(s) = index.as_ref() {
                         comps.push(s.text());
                         expr = Some(refr);
                     }
                 }
-                Some(Expr::Var(v)) => {
+                Expr::Var(v) => {
                     comps.push(v.text());
                     expr = None;
                 }
-                _ => bail!(format!("internal error: not a simplee ref {expr:?}")),
+                _ => bail!(e.span().error("invalid ref expression")),
             }
         }
         if let Some(d) = document {
diff --git a/src/utils.rs b/src/utils.rs
index 5d6c682e..45deb899 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -190,11 +190,12 @@ pub fn gather_functions(modules: &[Ref<Module>]) -> Result<FunctionTable> {
 }
 
 pub fn get_root_var(mut expr: &Expr) -> Result<SourceStr> {
+    let empty = expr.span().source_str().clone_empty();
     loop {
         match expr {
             Expr::Var(v) => return Ok(v.source_str()),
             Expr::RefDot { refr, .. } | Expr::RefBrack { refr, .. } => expr = refr,
-            _ => bail!("internal error: analyzer: could not get rule prefix"),
+            _ => return Ok(empty),
         }
     }
 }
diff --git a/tests/interpreter/cases/refr/tests.yaml b/tests/interpreter/cases/refr/tests.yaml
new file mode 100644
index 00000000..c5a48e80
--- /dev/null
+++ b/tests/interpreter/cases/refr/tests.yaml
@@ -0,0 +1,14 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+cases:
+  - note: ref
+    modules:
+      - |
+        package test
+        x = [k |
+          [1, 2, 3][[1, 2, 3][k]]
+        ]
+
+    query: data.test
+    want_result:
+      x: [0, 1]