Fix type back propagation for map keys

Issue: #3448 Both of these should be valid: ``` BEGIN { @stacks[@integer] = kstack; @integer = 1; } ``` And ``` BEGIN { @map1[@map2] = 1; @map2 = 1; for ($kv : @map1) { } } ```
bpftrace · Oct 24, 2024 · 673a6b9 · 673a6b9
1 parent cb5b1ea
commit 673a6b9
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 3 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -92,6 +92,8 @@ and this project adheres to
   - [#3517](https://github.com/bpftrace/bpftrace/pull/3517)
 - Fix verifier error when comparing result of len()
   - [#3308](https://github.com/bpftrace/bpftrace/issues/3308)
+- Fix type back propagation for map keys
+  - [#3536](https://github.com/bpftrace/bpftrace/pull/3536)
 #### Security
 #### Docs
 - Remove mention of unsupported character literals

diff --git a/src/ast/passes/semantic_analyser.cpp b/src/ast/passes/semantic_analyser.cpp
@@ -1662,8 +1662,10 @@ void SemanticAnalyser::validate_map_key(const SizedType &key,
 void SemanticAnalyser::visit(Map &map)
 {
   SizedType new_key_type = CreateNone();
+  bool key_is_map = false;
   if (map.key_expr) {
     Visit(map.key_expr);
+    key_is_map = map.key_expr->is_map;
     new_key_type = create_key_type(map.key_expr->type, map.key_expr->loc);
   }
 
@@ -1683,7 +1685,11 @@ void SemanticAnalyser::visit(Map &map)
         }
       }
     } else {
-      map_key_.insert({ map.ident, new_key_type });
+      // If the key used is a map, we might not have the type of it yet
+      // e.g. `BEGIN { @mymap[@i] = "hello"; @i = 1; }`
+      if (!key_is_map || !new_key_type.IsNoneTy()) {
+        map_key_.insert({ map.ident, new_key_type });
+      }
     }
   }
 
@@ -2422,8 +2428,10 @@ void SemanticAnalyser::visit(For &f)
     return;
 
   if (!mapkey || mapkey->IsNoneTy()) {
-    LOG(ERROR, map.loc, err_)
-        << "Maps used as for-loop expressions must have keys to iterate over";
+    if (is_final_pass()) {
+      LOG(ERROR, map.loc, err_)
+          << "Maps used as for-loop expressions must have keys to iterate over";
+    }
     return;
   }
 

diff --git a/tests/semantic_analyser.cpp b/tests/semantic_analyser.cpp
@@ -409,6 +409,8 @@ TEST(semantic_analyser, consistent_map_keys)
 {
   test("BEGIN { @x = 0; @x; }");
   test("BEGIN { @x[1] = 0; @x[2]; }");
+  test("BEGIN { @x[@y] = 5; @y = 1;}");
+  test("BEGIN { @x[@y[@z]] = 5; @y[2] = 1; @z = @x[0]; }");
 
   test_error("BEGIN { @x = 0; @x[1]; }", R"(
 stdin:1:17-22: ERROR: Argument mismatch for @x: trying to access with arguments: 'int64' when map expects no arguments
@@ -3872,6 +3874,8 @@ TEST(semantic_analyser, for_loop_map)
   test("BEGIN { @map[0] = 1; for ($kv : @map) { print($kv); } }");
   test("BEGIN { @map[0] = 1; for ($kv : @map) { print($kv.0); } }");
   test("BEGIN { @map[0] = 1; for ($kv : @map) { print($kv.1); } }");
+  test("BEGIN { @map1[@map2] = 1; @map2 = 1; for ($kv : @map1) { print($kv); } "
+       "}");
 }
 
 TEST(semantic_analyser, for_loop_map_declared_after)