Can now find the assignment node of a value

4 files changed, 35 insertions, 14 deletions

diff --git a/mesonbuild/ast/interpreter.py b/mesonbuild/ast/interpreter.py
index ce4b93c..b2cd3f5 100644
--- a/mesonbuild/ast/interpreter.py
+++ b/mesonbuild/ast/interpreter.py
@@ -15,12 +15,14 @@
 # This class contains the basic functionality needed to run any interpreter
 # or an interpreter-based tool.
 
+from .visitor import AstVisitor
 from .. import interpreterbase, mparser, mesonlib
 from .. import environment
 
 from ..interpreterbase import InvalidArguments, BreakRequest, ContinueRequest
 
 import os, sys
+from typing import List
 
 class DontCareObject(interpreterbase.InterpreterObject):
     pass
@@ -44,10 +46,12 @@ ADD_SOURCE = 0
 REMOVE_SOURCE = 1
 
 class AstInterpreter(interpreterbase.InterpreterBase):
-    def __init__(self, source_root, subdir):
+    def __init__(self, source_root: str, subdir: str, visitors: List[AstVisitor] = []):
         super().__init__(source_root, subdir)
+        self.visitors = visitors
         self.visited_subdirs = {}
         self.assignments = {}
+        self.reverse_assignment = {}
         self.funcs.update({'project': self.func_do_nothing,
                            'test': self.func_do_nothing,
                            'benchmark': self.func_do_nothing,
@@ -104,6 +108,11 @@ class AstInterpreter(interpreterbase.InterpreterBase):
     def func_do_nothing(self, node, args, kwargs):
         return True
 
+    def load_root_meson_file(self):
+        super().load_root_meson_file()
+        for i in self.visitors:
+            self.ast.accept(i)
+
     def func_subdir(self, node, args, kwargs):
         args = self.flatten_args(args)
         if len(args) != 1 or not isinstance(args[0], str):
@@ -134,6 +143,8 @@ class AstInterpreter(interpreterbase.InterpreterBase):
             raise me
 
         self.subdir = subdir
+        for i in self.visitors:
+            codeblock.accept(i)
         self.evaluate_codeblock(codeblock)
         self.subdir = prev_subdir
 
@@ -148,6 +159,8 @@ class AstInterpreter(interpreterbase.InterpreterBase):
         if node.var_name not in self.assignments:
             self.assignments[node.var_name] = []
         self.assignments[node.var_name] += [node.value] # Save a reference to the value node
+        if hasattr(node.value, 'ast_id'):
+            self.reverse_assignment[node.value.ast_id] = node
         self.evaluate_statement(node.value) # Evaluate the value just in case
 
     def evaluate_indexing(self, node):
@@ -185,6 +198,8 @@ class AstInterpreter(interpreterbase.InterpreterBase):
     def assignment(self, node):
         assert(isinstance(node, mparser.AssignmentNode))
         self.assignments[node.var_name] = [node.value] # Save a reference to the value node
+        if hasattr(node.value, 'ast_id'):
+            self.reverse_assignment[node.value.ast_id] = node
         self.evaluate_statement(node.value) # Evaluate the value just in case
 
     def flatten_args(self, args, include_unknown_args: bool = False):
diff --git a/mesonbuild/ast/introspection.py b/mesonbuild/ast/introspection.py
index b6a523b..0917015 100644
--- a/mesonbuild/ast/introspection.py
+++ b/mesonbuild/ast/introspection.py
@@ -34,8 +34,8 @@ class IntrospectionHelper:
 class IntrospectionInterpreter(AstInterpreter):
     # Interpreter to detect the options without a build directory
     # Most of the code is stolen from interperter.Interpreter
-    def __init__(self, source_root, subdir, backend, cross_file=None, subproject='', subproject_dir='subprojects', env=None):
-        super().__init__(source_root, subdir)
+    def __init__(self, source_root, subdir, backend, visitors=[], cross_file=None, subproject='', subproject_dir='subprojects', env=None):
+        super().__init__(source_root, subdir, visitors=visitors)
 
         options = IntrospectionHelper(cross_file)
         self.cross_file = cross_file
@@ -162,9 +162,9 @@ class IntrospectionInterpreter(AstInterpreter):
                 # Try to resolve the ID and append the node to the queue
                 id = curr.value
                 if id in self.assignments and self.assignments[id]:
-                    node = self.assignments[id][0]
-                    if isinstance(node, (mparser.ArrayNode, mparser.IdNode, mparser.FunctionNode)):
-                        srcqueue += [node]
+                    tmp_node = self.assignments[id][0]
+                    if isinstance(tmp_node, (mparser.ArrayNode, mparser.IdNode, mparser.FunctionNode)):
+                        srcqueue += [tmp_node]
             if arg_node is None:
                 continue
             elemetary_nodes = list(filter(lambda x: isinstance(x, (str, mparser.StringNode)), arg_node.arguments))
diff --git a/mesonbuild/mparser.py b/mesonbuild/mparser.py
index ed0dc1b..942a73f 100644
--- a/mesonbuild/mparser.py
+++ b/mesonbuild/mparser.py
@@ -358,7 +358,8 @@ class FunctionNode(BaseNode):
         self.args = args
 
 class AssignmentNode(BaseNode):
-    def __init__(self, lineno, colno, var_name, value):
+    def __init__(self, subdir, lineno, colno, var_name, value):
+        self.subdir = subdir
         self.lineno = lineno
         self.colno = colno
         self.var_name = var_name
@@ -366,7 +367,8 @@ class AssignmentNode(BaseNode):
         self.value = value
 
 class PlusAssignmentNode(BaseNode):
-    def __init__(self, lineno, colno, var_name, value):
+    def __init__(self, subdir, lineno, colno, var_name, value):
+        self.subdir = subdir
         self.lineno = lineno
         self.colno = colno
         self.var_name = var_name
@@ -522,13 +524,13 @@ class Parser:
             value = self.e1()
             if not isinstance(left, IdNode):
                 raise ParseException('Plusassignment target must be an id.', self.getline(), left.lineno, left.colno)
-            return PlusAssignmentNode(left.lineno, left.colno, left.value, value)
+            return PlusAssignmentNode(left.subdir,left.lineno, left.colno, left.value, value)
         elif self.accept('assign'):
             value = self.e1()
             if not isinstance(left, IdNode):
                 raise ParseException('Assignment target must be an id.',
                                      self.getline(), left.lineno, left.colno)
-            return AssignmentNode(left.lineno, left.colno, left.value, value)
+            return AssignmentNode(left.subdir, left.lineno, left.colno, left.value, value)
         elif self.accept('questionmark'):
             if self.in_ternary:
                 raise ParseException('Nested ternary operators are not allowed.',
diff --git a/mesonbuild/rewriter.py b/mesonbuild/rewriter.py
index 60c762e..6752189 100644
--- a/mesonbuild/rewriter.py
+++ b/mesonbuild/rewriter.py
@@ -292,9 +292,10 @@ rewriter_func_kwargs = {
 class Rewriter:
     def __init__(self, sourcedir: str, generator: str = 'ninja'):
         self.sourcedir = sourcedir
-        self.interpreter = IntrospectionInterpreter(sourcedir, '', generator)
-        self.id_generator = AstIDGenerator()
+        self.interpreter = IntrospectionInterpreter(sourcedir, '', generator, visitors = [AstIDGenerator(), AstIndentationGenerator()])
         self.modefied_nodes = []
+        self.to_remove_nodes = []
+        self.to_add_nodes = []
         self.functions = {
             'kwargs': self.process_kwargs,
             'target': self.process_target,
@@ -306,8 +307,6 @@ class Rewriter:
         self.interpreter.analyze()
         mlog.log('  -- Project:', mlog.bold(self.interpreter.project_data['descriptive_name']))
         mlog.log('  -- Version:', mlog.cyan(self.interpreter.project_data['version']))
-        self.interpreter.ast.accept(AstIndentationGenerator())
-        self.interpreter.ast.accept(self.id_generator)
 
     def add_info(self, cmd_type: str, cmd_id: str, data: dict):
         if self.info_dump is None:
@@ -456,6 +455,11 @@ class Rewriter:
         if num_changed > 0 and node not in self.modefied_nodes:
             self.modefied_nodes += [node]
 
+    def find_assignment_node(self, node: mparser.BaseNode) -> mparser.AssignmentNode:
+        if hasattr(node, 'ast_id') and node.ast_id in self.interpreter.reverse_assignment:
+            return self.interpreter.reverse_assignment[node.ast_id]
+        return None
+
     @RequiredKeys(rewriter_keys['target'])
     def process_target(self, cmd):
         mlog.log('Processing target', mlog.bold(cmd['target']), 'operation', mlog.cyan(cmd['operation']))