From 1c8bb79e147e47c7c6658fbc9fb4b0aaceff7a87 Mon Sep 17 00:00:00 2001
From: thsparks <thsparks@microsoft.com>
Date: Wed, 1 May 2024 11:03:14 -0700
Subject: [PATCH] Have validator return Map of passingVarDefinitions instead of
 a list and a count. For now just flatten it in runValidatorPlan itself to
 prep for recursive calls...

---
 pxteditor/code-validation/runValidatorPlan.ts |  9 +++++++-
 .../code-validation/validateVariableUsage.ts  | 21 +++++++------------
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/pxteditor/code-validation/runValidatorPlan.ts b/pxteditor/code-validation/runValidatorPlan.ts
index fa4bf8b93c3a..1022341fd0c9 100644
--- a/pxteditor/code-validation/runValidatorPlan.ts
+++ b/pxteditor/code-validation/runValidatorPlan.ts
@@ -115,5 +115,12 @@ function runVariableUsageValidation(usedBlocks: Blockly.Block[], inputs: pxt.blo
         count: inputs.count,
         name: inputs.name
     });
-    return [blockResults.passingVarDefinitions, blockResults.passed];
+
+    // Flatten the map of passing variable definition blocks
+    const passingVarDefinitions: Blockly.Block[] = [];
+    for (const blocks of blockResults.passingVarDefinitions.values()) {
+        passingVarDefinitions.push(...blocks);
+    }
+
+    return [passingVarDefinitions, blockResults.passed];
 }
diff --git a/pxteditor/code-validation/validateVariableUsage.ts b/pxteditor/code-validation/validateVariableUsage.ts
index e88cbcbe793a..97c023f41597 100644
--- a/pxteditor/code-validation/validateVariableUsage.ts
+++ b/pxteditor/code-validation/validateVariableUsage.ts
@@ -13,12 +13,11 @@ export function validateVariableUsage({
     count: number;
     name?: String;
 }): {
-    passingVarDefinitions: Blockly.Block[];
-    passingVarCount: number;
+    passingVarDefinitions: Map<string, Blockly.Block[]>;
     passed: boolean;
 } {
     const varDefinitionBlocks: Map<string, Blockly.Block[]> = new Map();
-    const usedVars: Set<string> = new Set(); // Use set so we don't double count vars
+    const usedVars: Set<string> = new Set();
 
     for (const block of usedBlocks) {
         if (!block.isEnabled()) {
@@ -45,18 +44,12 @@ export function validateVariableUsage({
 
     // Var passes check if it is both used and defined.
     // We return the definition blocks to allow for recursively checking how the var was set.
-    // Track passingVarCount separately, since a var could have more than one set block.
-    const passingVarDefinitions: Blockly.Block[] = [];
-    let passingVarCount = 0;
-    for (const varName of usedVars) {
-        const varBlocks = varDefinitionBlocks.get(varName) ?? [];
-        if (varBlocks.length > 0) {
-            passingVarCount++;
-            for (const varBlock of varBlocks) {
-                passingVarDefinitions.push(varBlock);
-            }
+    const passingVarDefinitions = new Map<string, Blockly.Block[]>();
+    for (const [varName, definitionBlocks] of varDefinitionBlocks) {
+        if (usedVars.has(varName)) {
+            passingVarDefinitions.set(varName, definitionBlocks);
         }
     }
 
-    return { passingVarDefinitions, passingVarCount, passed: passingVarCount >= count };
+    return { passingVarDefinitions, passed: passingVarDefinitions.size >= count };
 }