Analysis: Follow state into the locals it is assigned to

javascript/packages/analysis/src/dependency-index.ts

@@ -38,6 +38,39 @@ function referencesAny(code: string | undefined, aliases: string[]): boolean {
   return aliases.some(alias => referencesState(code, alias))
 }
 
+const ASSIGNMENT_NODES = new Set([
+  "LocalVariableWriteNode",
+  "LocalVariableOrWriteNode",
+  "LocalVariableAndWriteNode",
+  "LocalVariableOperatorWriteNode",
+])
+
+function assignedNames(node: Node, source: string, aliases: string[]): string[] {
+  const prism = (node as { prismNode?: unknown }).prismNode as any
+
+  if (!prism) return []
+
+  const names: string[] = []
+
+  const walk = (candidate: any) => {
+    if (!candidate) return
+
+    if (ASSIGNMENT_NODES.has(candidate.constructor?.name) && candidate.value?.location) {
+      const { startOffset, length } = candidate.value.location
+      const right = source.slice(startOffset, startOffset + length)
+      const name = String(candidate.name)
+
+      if (referencesAny(right, aliases) && !names.includes(name)) names.push(name)
+    }
+
+    for (const child of candidate.compactChildNodes?.() ?? []) walk(child)
+  }
+
+  walk(prism)
+
+  return names
+}
+
 function blockBindings(node: Node, aliases: string[]): string[] {
   if (!isERBBlockNode(node)) return []
   if (!referencesAny(expressionOf(node), aliases)) return []
@@ -164,6 +197,13 @@ export function affectedNodes(backend: HerbBackend, source: string, state: strin
       record(node, "expression", expressionOf(node))
     }
 
+    if (isERBContentNode(node)) {
+      for (const name of assignedNames(node, source, aliases)) {
+        if (!aliases.includes(name)) aliases.push(name)
+      }
+    }
+
+    const outer = [...aliases]
     const bound = blockBindings(node, aliases)
 
     aliases.push(...bound)
@@ -174,7 +214,7 @@ export function affectedNodes(backend: HerbBackend, source: string, state: strin
       path.pop()
     }
 
-    for (const name of bound) aliases.splice(aliases.indexOf(name), 1)
+    if (bound.length > 0) aliases.splice(0, aliases.length, ...outer)
   }
 
   for (const [index, child] of childrenOf(document).entries()) {
@@ -190,7 +230,7 @@ export function dependencyIndex(backend: HerbBackend, file: string, source: stri
   const dependencies = collectTemplateDependencies(backend, file, source, options)
   const index: Record<string, AffectedNode[]> = {}
 
-  for (const state of [...dependencies.instanceVariables, ...dependencies.constants]) {
+  for (const state of [...dependencies.instanceVariables, ...dependencies.constants, ...dependencies.localsDeclared]) {
     const nodes = affectedNodes(backend, source, state)
 
     if (nodes.length > 0) index[state] = nodes

javascript/packages/analysis/test/dependency-index.test.ts

@@ -15,6 +15,34 @@ describe("dependencyIndex", () => {
     return dependencyIndex(Herb, FILE, source)
   }
 
+  test("follows state into a local assigned from it", () => {
+    const nodes = affectedNodes(Herb, `<% total = @items.size %><p><%= total %></p>`, "@items")
+
+    expect(nodes.map(node => node.expression)).toEqual(["total = @items.size", "total"])
+  })
+
+  test("follows state through a chain of assignments", () => {
+    const nodes = affectedNodes(Herb, `<% a = @items.size %><% b = a * 2 %><p><%= b %></p>`, "@items")
+
+    expect(nodes.map(node => node.expression)).toEqual(["a = @items.size", "b = a * 2", "b"])
+  })
+
+  test("leaves a local assigned from something else", () => {
+    expect(affectedNodes(Herb, `<% other = 5 %><p><%= other %></p>`, "@items")).toEqual([])
+  })
+
+  test("does not take a comparison for an assignment", () => {
+    const nodes = affectedNodes(Herb, `<% if @items == other %><p><%= other %></p><% end %>`, "@items")
+
+    expect(nodes.map(node => node.expression)).toEqual(["if @items == other"])
+  })
+
+  test("stops a local assigned inside a block at the end of it", () => {
+    const nodes = affectedNodes(Herb, `<% @rows.each do |r| %><% inner = r.x %><%= inner %><% end %><%= inner %>`, "@rows")
+
+    expect(nodes.map(node => node.expression)).toEqual(["@rows.each do |r|", "inner = r.x", "inner"])
+  })
+
   test("follows state through a block parameter", () => {
     const nodes = affectedNodes(Herb, `<ul><% @items.each do |item| %><li><%= item.name %></li><% end %></ul>`, "@items")
 

lib/herb/analysis/template_dependencies.rb

@@ -95,7 +95,7 @@ def dependency_index(file_path)
 
         index = {} #: Hash[String, Array[Hash[Symbol, untyped]]]
 
-        (result.instance_variables + result.constants).each do |state|
+        (result.instance_variables + result.constants + result.locals_declared).each do |state|
           nodes = affected_nodes(file_path, state)
           index[state] = nodes if nodes.any?
         end

lib/herb/analysis/template_dependencies/node_dependency_collector.rb

@@ -8,6 +8,12 @@ class TemplateDependencies
       class NodeDependencyCollector < ::Herb::Visitor
         BRANCH_BODY_PROPERTIES = [:statements, :body, :children, :conditions].freeze #: Array[Symbol]
         BRANCH_CONTINUATION_PROPERTIES = [:subsequent, :else_clause, :rescue_clause, :ensure_clause].freeze #: Array[Symbol]
+        ASSIGNMENT_NODES = [
+          Prism::LocalVariableWriteNode,
+          Prism::LocalVariableOrWriteNode,
+          Prism::LocalVariableAndWriteNode,
+          Prism::LocalVariableOperatorWriteNode,
+        ].freeze #: Array[untyped]
 
         attr_reader :affected
 
@@ -42,6 +48,8 @@ def visit_html_open_tag_node(node)
 
         def visit_erb_content_node(node)
           check_erb_expression(node, :text_content)
+
+          bind_assignments(node)
         end
 
         def visit_erb_if_node(node)
@@ -125,8 +133,9 @@ def check_block_for_state(node, type)
         end
 
         def visit_branching_node(node)
-          bound = bindings_for(node)
-          @aliases.merge(bound)
+          outer = @aliases.dup
+
+          @aliases.merge(bindings_for(node))
 
           BRANCH_BODY_PROPERTIES.each do |property|
             next unless node.respond_to?(property)
@@ -143,7 +152,44 @@ def visit_branching_node(node)
             visit(child)
           end
         ensure
-          @aliases.subtract(bound) if bound
+          @aliases.replace(outer) if outer
+        end
+
+        def bind_assignments(node)
+          code = node.content&.value&.strip
+
+          return unless code
+
+          assigned_names(code).each { |name| @aliases.add(name) }
+        end
+
+        def assigned_names(code)
+          result = Prism.parse(code)
+
+          return [] if result.errors.any?
+
+          names = [] #: Array[String]
+          collect_assignments(result.value, names)
+
+          names
+        end
+
+        def collect_assignments(node, names)
+          if ASSIGNMENT_NODES.any? { |type| node.is_a?(type) }
+            right = right_hand_side(node)
+
+            names << node.name.to_s if right && references_state?(right)
+          end
+
+          node.child_nodes.compact.each { |child| collect_assignments(child, names) }
+        end
+
+        def right_hand_side(node)
+          return nil unless node.respond_to?(:value)
+
+          value = node.value #: untyped
+
+          value&.slice
         end
 
         def bindings_for(node)

rust/herb-analysis/src/state_flow.rs

@@ -104,7 +104,7 @@ impl StateFlow {
 
     for (index, child) in result.value.children.iter().enumerate() {
       path.push(index);
-      collect_affected(child, &mut aliases, &mut path, &mut affected);
+      collect_affected(child, &source, &mut aliases, &mut path, &mut affected);
       path.pop();
     }
 
@@ -115,7 +115,12 @@ impl StateFlow {
     let result = self.analyze(file);
     let mut index = BTreeMap::new();
 
-    for state in result.instance_variables.iter().chain(result.constants.iter()) {
+    for state in result
+      .instance_variables
+      .iter()
+      .chain(result.constants.iter())
+      .chain(result.locals_declared.iter())
+    {
       let nodes = self.affected_nodes(file, state);
 
       if !nodes.is_empty() {
@@ -295,7 +300,7 @@ fn is_word_byte(byte: u8) -> bool {
   byte.is_ascii_alphanumeric() || byte == b'_'
 }
 
-fn collect_affected(node: &AnyNode, aliases: &mut Vec<String>, path: &mut Vec<usize>, affected: &mut Vec<AffectedNode>) {
+fn collect_affected(node: &AnyNode, source: &str, aliases: &mut Vec<String>, path: &mut Vec<usize>, affected: &mut Vec<AffectedNode>) {
   let kind = match node {
     AnyNode::ERBContentNode(_) => Some("text_content"),
     AnyNode::ERBIfNode(_) => Some("conditional"),
@@ -326,17 +331,25 @@ fn collect_affected(node: &AnyNode, aliases: &mut Vec<String>, path: &mut Vec<us
     collect_attributes(element, aliases, path, affected);
   }
 
-  let bound = block_bindings(node, aliases);
+  for name in assigned_names(node, source, aliases) {
+    if !aliases.contains(&name) {
+      aliases.push(name);
+    }
+  }
 
-  aliases.extend(bound.iter().cloned());
+  let outer = aliases.clone();
+
+  aliases.extend(block_bindings(node, aliases));
 
   for (index, child) in any_children(node).into_iter().enumerate() {
     path.push(index);
-    collect_affected(child, aliases, path, affected);
+    collect_affected(child, source, aliases, path, affected);
     path.pop();
   }
 
-  aliases.retain(|name| !bound.contains(name));
+  if aliases.len() != outer.len() {
+    *aliases = outer;
+  }
 }
 
 fn is_block(node: &AnyNode) -> bool {
@@ -392,6 +405,50 @@ fn references_any(code: &str, aliases: &[String]) -> bool {
   aliases.iter().any(|name| expression_references(code, name))
 }
 
+const ASSIGNMENT_NODES: [&str; 4] = [
+  "LocalVariableWriteNode",
+  "LocalVariableOrWriteNode",
+  "LocalVariableAndWriteNode",
+  "LocalVariableOperatorWriteNode",
+];
+
+fn assigned_names(node: &AnyNode, source: &str, aliases: &[String]) -> Vec<String> {
+  let AnyNode::ERBContentNode(inner) = node else {
+    return Vec::new();
+  };
+
+  let Some(prism) = inner.prism() else {
+    return Vec::new();
+  };
+
+  let mut names = Vec::new();
+
+  collect_assignments(prism, source, aliases, &mut names);
+
+  names
+}
+
+fn collect_assignments(node: &herb::prism::PrismNode, source: &str, aliases: &[String], names: &mut Vec<String>) {
+  if ASSIGNMENT_NODES.contains(&node.node_type.as_str()) {
+    if let Some(name) = node.name.as_ref() {
+      let assigns = node.children.iter().any(|child| {
+        source
+          .get(child.start_offset..child.end_offset)
+          .map(|right| references_any(right, aliases))
+          .unwrap_or(false)
+      });
+
+      if assigns && !names.contains(name) {
+        names.push(name.clone());
+      }
+    }
+  }
+
+  for child in &node.children {
+    collect_assignments(child, source, aliases, names);
+  }
+}
+
 fn collect_attributes(element: &herb::nodes::HTMLElementNode, aliases: &[String], path: &[usize], affected: &mut Vec<AffectedNode>) {
   let Some(open_tag) = element.open_tag.as_ref() else {
     return;

rust/herb-analysis/tests/state_flow_test.rs

@@ -348,3 +348,56 @@ fn reports_what_the_ruby_collector_reports() {
     reported
   );
 }
+
+fn expressions_for(name: &str, template: &str, state: &str) -> Vec<String> {
+  let project = Project::new(name);
+  let entry = project.write("app/views/posts/index.html.erb", template);
+
+  project
+    .flow()
+    .affected_nodes(&entry, state)
+    .into_iter()
+    .filter_map(|node| node.expression)
+    .collect()
+}
+
+#[test]
+fn follows_state_into_a_local_assigned_from_it() {
+  assert_eq!(
+    vec!["total = @items.size".to_string(), "total".to_string()],
+    expressions_for("assign_simple", "<% total = @items.size %><p><%= total %></p>", "@items")
+  );
+}
+
+#[test]
+fn follows_state_through_a_chain_of_assignments() {
+  assert_eq!(
+    vec!["a = @items.size".to_string(), "b = a * 2".to_string(), "b".to_string()],
+    expressions_for("assign_chain", "<% a = @items.size %><% b = a * 2 %><p><%= b %></p>", "@items")
+  );
+}
+
+#[test]
+fn leaves_a_local_assigned_from_something_else() {
+  assert!(expressions_for("assign_unrelated", "<% other = 5 %><p><%= other %></p>", "@items").is_empty());
+}
+
+#[test]
+fn does_not_take_a_comparison_for_an_assignment() {
+  assert_eq!(
+    vec!["if @items == other".to_string()],
+    expressions_for("assign_comparison", "<% if @items == other %><p><%= other %></p><% end %>", "@items")
+  );
+}
+
+#[test]
+fn stops_a_local_assigned_inside_a_block_at_the_end_of_it() {
+  assert_eq!(
+    vec!["@rows.each do |r|".to_string(), "inner = r.x".to_string(), "inner".to_string()],
+    expressions_for(
+      "assign_scope",
+      "<% @rows.each do |r| %><% inner = r.x %><%= inner %><% end %><%= inner %>",
+      "@rows"
+    )
+  );
+}