Remove AstVarScope && AstVarEnv (#60373)

DrRyanHuang · web-flow · commit 3055fac26c5c · 2023-12-27T00:21:17.000+08:00
diff --git a/python/paddle/jit/dy2static/static_analysis.py b/python/paddle/jit/dy2static/static_analysis.py
@@ -40,110 +40,6 @@ def __init__(self, node):
         self.node_var_type = {NodeVarType.UNKNOWN}
 
 
-class AstVarScope:
-    """
-    AstVarScope is a class holding the map from current scope variable to its
-    type.
-    """
-
-    SCOPE_TYPE_SCRIPT = 0
-    SCOPE_TYPE_FUNCTION = 1
-    SCOPE_TYPE_CLASS = 2
-
-    def __init__(
-        self, scope_name='', scope_type=SCOPE_TYPE_SCRIPT, parent_scope=None
-    ):
-        self.sub_scopes = []
-        self.name_to_id = {}
-        self.id_to_type = {}
-        self.cur_id = 0
-
-        self.scope_name = scope_name
-        self.scope_type = scope_type
-        self.parent_scope = parent_scope
-        if parent_scope is not None:
-            parent_scope.sub_scopes.append(self)
-
-    def add_var_type(self, var_name, node_var_type):
-        var_type = self.get_var_type(var_name)
-        if var_type == {NodeVarType.UNKNOWN}:
-            self.set_var_type(var_name, node_var_type)
-        else:
-            if isinstance(node_var_type, set):
-                var_type.update(node_var_type)
-            else:
-                var_type.add(node_var_type)
-
-    def set_var_type(self, var_name, node_var_type):
-        if var_name in self.name_to_id:
-            num_id = self.name_to_id[var_name]
-        else:
-            num_id = self.cur_id
-            self.cur_id += 1
-            self.name_to_id[var_name] = num_id
-        self.id_to_type[num_id] = (
-            node_var_type if isinstance(node_var_type, set) else {node_var_type}
-        )
-
-    def get_var_type(self, var_name):
-        if var_name in self.name_to_id:
-            num_id = self.name_to_id[var_name]
-            return self.id_to_type[num_id]
-        if self.parent_scope is None:
-            return {NodeVarType.UNKNOWN}
-        return self.parent_scope.get_var_type(var_name)
-
-
-class AstVarEnv:
-    """
-    A class maintains scopes and mapping from name strings to type.
-    """
-
-    def __init__(self):
-        self.cur_scope = AstVarScope()
-
-    def enter_scope(self, scope_name, scope_type):
-        self.cur_scope = AstVarScope(
-            scope_name, scope_type, parent_scope=self.cur_scope
-        )
-        return self.cur_scope
-
-    def exit_scope(self):
-        assert self.cur_scope.parent_scope is not None, (
-            "Call exit_scope in "
-            "AstVarEnv when current scope doesn't have parent scope."
-        )
-        self.cur_scope = self.cur_scope.parent_scope
-        return self.cur_scope
-
-    def get_parent_scope(self):
-        assert self.cur_scope.parent_scope is not None, (
-            "Call parent_scope in "
-            "AstVarEnv when current scope doesn't have parent scope."
-        )
-        return self.cur_scope.parent_scope
-
-    def add_var_type(self, var_name, node_var_type):
-        self.cur_scope.add_var_type(var_name, node_var_type)
-
-    def set_var_type(self, var_name, node_var_type):
-        self.cur_scope.set_var_type(var_name, node_var_type)
-
-    def get_var_type(self, var_name):
-        return self.cur_scope.get_var_type(var_name)
-
-    def get_scope_var_type(self):
-        '''
-        Returns a dict mapping from variable name to type. Used for debug and
-        test.
-        '''
-        cur_scope_dict = {}
-        for name in self.cur_scope.name_to_id:
-            node_var_type = self.cur_scope.get_var_type(name)
-            cur_scope_dict[name] = node_var_type
-        return cur_scope_dict
-
-
 class StaticAnalysisVisitor:
     """
     A class that does static analysis
@@ -157,7 +53,6 @@ def run(self, ast_root):
         self.node_wrapper_root = None
         self.ancestor_wrappers = []
         self.node_to_wrapper_map = {}
-        self.var_env = AstVarEnv()
 
         self.dfs_visit(ast_root)
 
@@ -179,16 +74,7 @@ def dfs_visit(self, node):
 
         self.ancestor_wrappers.append(cur_wrapper)
         for child in gast.iter_child_nodes(node):
-            if isinstance(child, (gast.FunctionDef, gast.AsyncFunctionDef)):
-                # TODO: current version is function name mapping to its type
-                # consider complex case involving parameters
-                self.var_env.enter_scope(
-                    child.name, AstVarScope.SCOPE_TYPE_FUNCTION
-                )
-                func_type = self.dfs_visit(child)
-                self.var_env.exit_scope()
-            else:
-                self.dfs_visit(child)
+            self.dfs_visit(child)
         self.ancestor_wrappers.pop()
 
         cur_wrapper.node_var_type = self._get_node_var_type(cur_wrapper)
@@ -200,9 +86,6 @@ def get_node_wrapper_root(self):
     def get_node_to_wrapper_map(self):
         return self.node_to_wrapper_map
 
-    def get_var_env(self):
-        return self.var_env
-
     def is_tensor_node(self, node):
         tensor_types = {NodeVarType.TENSOR, NodeVarType.PADDLE_RETURN_TYPES}
         node_wrapper = self.node_to_wrapper_map.get(node, None)
@@ -262,15 +145,13 @@ def _get_node_var_type(self, cur_wrapper):
             for target in node.targets:
                 if isinstance(target, gast.Name):
                     self.node_to_wrapper_map[target].node_var_type = ret_type
-                    self.var_env.set_var_type(target.id, ret_type)
                 # Handle statements like `a, b = paddle.shape(x)`
                 elif isinstance(target, gast.Tuple):
                     for sub_target in target.elts:
                         if isinstance(sub_target, gast.Name):
                             self.node_to_wrapper_map[
                                 sub_target
                             ].node_var_type = ret_type
-                            self.var_env.set_var_type(sub_target.id, ret_type)
             return ret_type
 
         if isinstance(node, gast.AnnAssign):
@@ -289,7 +170,6 @@ def _get_node_var_type(self, cur_wrapper):
                     ret_type = node_value_type
             if isinstance(node.target, gast.Name):
                 self.node_to_wrapper_map[node.target].node_var_type = ret_type
-                self.var_env.set_var_type(node.target.id, ret_type)
             return ret_type
 
         if isinstance(node, gast.Name):
@@ -304,22 +184,14 @@ def _get_node_var_type(self, cur_wrapper):
             ):
                 return self._get_func_argument_type(parent_node_wrapper, node)
 
-            return self.var_env.get_var_type(node.id)
+            return {NodeVarType.UNKNOWN}
 
         if isinstance(node, gast.Return):
             # If return nothing:
             if node.value is None:
                 return {NodeVarType.NONE}
 
-            return_type = self.node_to_wrapper_map[node.value].node_var_type
-            assert (
-                self.var_env.cur_scope.scope_type
-                == AstVarScope.SCOPE_TYPE_FUNCTION
-            ), "Return at non-function scope"
-            func_name = self.var_env.cur_scope.scope_name
-            parent_scope = self.var_env.get_parent_scope()
-            parent_scope.add_var_type(func_name, return_type)
-            return return_type
+            return {NodeVarType.UNKNOWN}
 
         if isinstance(node, gast.Call):
             if is_dygraph_api(node):
@@ -333,7 +205,7 @@ def _get_node_var_type(self, cur_wrapper):
                 return {NodeVarType.NUMPY_NDARRAY}
 
             if isinstance(node.func, gast.Name):
-                return self.var_env.get_var_type(node.func.id)
+                return {NodeVarType.UNKNOWN}
         if isinstance(node, gast.Subscript):
             if self.is_tensor_node(node.value):
                 return {NodeVarType.TENSOR}
@@ -363,7 +235,6 @@ def _get_func_argument_type(self, parent_node_wrapper, node):
         var_type = {NodeVarType.UNKNOWN}
         if node.annotation is not None:
             var_type = {NodeVarType.type_from_annotation(node.annotation)}
-            self.var_env.set_var_type(node.id, var_type)
 
         # if annotation and value(Constant) are diffent type, we use value type
         if parent_node.defaults:
@@ -374,7 +245,4 @@ def _get_func_argument_type(self, parent_node_wrapper, node):
                 if isinstance(defaults_node, gast.Constant):
                     var_type = self._get_constant_node_type(defaults_node)
 
-                    # Add node with identified type into cur_env.
-                    self.var_env.set_var_type(node.id, var_type)
-
         return var_type
diff --git a/test/dygraph_to_static/test_static_analysis.py b/test/dygraph_to_static/test_static_analysis.py
@@ -200,27 +200,6 @@ def test_construct_node_wrapper(self):
             node_to_wrapper_map = visitor.get_node_to_wrapper_map()
             self._check_wrapper(wrapper_root, node_to_wrapper_map)
 
-    def test_var_env(self):
-        for i, func in enumerate(test_funcs):
-            var_type = result_var_type[i]
-            test_source_code = inspect.getsource(func)
-            ast_root = gast.parse(test_source_code)
-            print(gast.dump(ast_root))
-            visitor = StaticAnalysisVisitor(ast_root)
-            var_env = visitor.get_var_env()
-
-            # There must be 1 sub scope for the test function
-            self.assertEqual(1, len(var_env.cur_scope.sub_scopes))
-            var_env.cur_scope = var_env.cur_scope.sub_scopes[0]
-
-            scope_var_type = var_env.get_scope_var_type()
-            print(scope_var_type)
-            self.assertEqual(len(scope_var_type), len(var_type))
-            for name in scope_var_type:
-                print("Test var name %s" % (name))
-                self.assertTrue(name in var_type)
-                self.assertEqual(scope_var_type[name], var_type[name])
-
 
 if __name__ == '__main__':
     unittest.main()
