Basic locator

Author: youknowone

ast/Cargo.toml

@@ -8,8 +8,9 @@ repository = "https://github.com/RustPython/RustPython"
 license = "MIT"
 
 [features]
-default = ["constant-optimization", "fold"]
+default = ["constant-optimization", "fold", "location"]
 constant-optimization = ["fold"]
+location = []
 fold = []
 unparse = ["rustpython-literal"]
 

ast/asdl_rs.py

@@ -65,8 +65,9 @@ def asdl_of(name, obj):
 class EmitVisitor(asdl.VisitorBase):
     """Visit that emits lines"""
 
-    def __init__(self, file):
+    def __init__(self, file, typeinfo):
         self.file = file
+        self.typeinfo = typeinfo
         self.identifiers = set()
         super(EmitVisitor, self).__init__()
 
@@ -163,10 +164,13 @@ def rust_field(field_name):
         return field_name
 
 
+def product_has_expr(product):
+    return any(f.type != "identifier" for f in product.fields)
+
+
 class TypeInfoEmitVisitor(EmitVisitor):
     def __init__(self, file, typeinfo):
-        self.typeinfo = typeinfo
-        super().__init__(file)
+        super().__init__(file, typeinfo)
 
     def has_userdata(self, typ):
         return self.typeinfo[typ].has_userdata
@@ -300,7 +304,7 @@ def visitProduct(self, product, name, depth):
         if product.attributes:
             dataname = rustname + "Data"
         self.emit_attrs(depth)
-        has_expr = any(f.type != "identifier" for f in product.fields)
+        has_expr = product_has_expr(product)
         if has_expr:
             datadef = f"{dataname}{generics}"
         else:
@@ -327,7 +331,11 @@ def visitModule(self, mod, depth):
         self.emit("type Error;", depth + 1)
         self.emit(
             "fn map_user(&mut self, user: U) -> Result<Self::TargetU, Self::Error>;",
-            depth + 2,
+            depth + 1,
+        )
+        self.emit(
+            "fn map_located<T>(&mut self, located: Located<T, U>) -> Result<Located<T, Self::TargetU>, Self::Error> { let custom = self.map_user(located.custom)?; Ok(Located { range: located.range, custom, node: located.node }) }",
+            depth + 1,
         )
         for dfn in mod.dfns:
             self.visit(dfn, depth + 2)
@@ -352,7 +360,7 @@ def visitModule(self, mod, depth):
             depth,
         )
         self.emit(
-            "Ok(Located { custom: folder.map_user(node.custom)?, range: node.range, node: f(folder, node.node)? })",
+            "let node = folder.map_located(node)?; Ok(Located { custom: node.custom, range: node.range, node: f(folder, node.node)? })",
             depth + 1,
         )
         self.emit("}", depth)
@@ -575,11 +583,15 @@ def visitSum(self, sum, name, depth):
         rustname = enumname = get_rust_type(name)
         if sum.attributes:
             rustname = enumname + "Kind"
+        if is_simple(sum) or not self.typeinfo[name].has_userdata:
+            custom = ""
+        else:
+            custom = "<LocationRange>"
 
-        self.emit(f"impl NamedNode for ast::{rustname} {{", depth)
+        self.emit(f"impl NamedNode for ast::{rustname}{custom} {{", depth)
         self.emit(f"const NAME: &'static str = {json.dumps(name)};", depth + 1)
         self.emit("}", depth)
-        self.emit(f"impl Node for ast::{rustname} {{", depth)
+        self.emit(f"impl Node for ast::{rustname}{custom} {{", depth)
         self.emit(
             "fn ast_to_object(self, _vm: &VirtualMachine) -> PyObjectRef {", depth + 1
         )
@@ -609,11 +621,16 @@ def visitProduct(self, product, name, depth):
         structname = get_rust_type(name)
         if product.attributes:
             structname += "Data"
+            custom = ""
+        if product_has_expr(product):
+            custom = "<LocationRange>"
+        else:
+            custom = ""
 
-        self.emit(f"impl NamedNode for ast::{structname} {{", depth)
+        self.emit(f"impl NamedNode for ast::{structname}{custom} {{", depth)
         self.emit(f"const NAME: &'static str = {json.dumps(name)};", depth + 1)
         self.emit("}", depth)
-        self.emit(f"impl Node for ast::{structname} {{", depth)
+        self.emit(f"impl Node for ast::{structname}{custom} {{", depth)
         self.emit(
             "fn ast_to_object(self, _vm: &VirtualMachine) -> PyObjectRef {", depth + 1
         )
@@ -656,7 +673,14 @@ def gen_sum_fromobj(self, sum, sumname, enumname, rustname, depth):
         for cons in sum.types:
             self.emit(f"if _cls.is(Node{cons.name}::static_type()) {{", depth)
             if cons.fields:
-                self.emit(f"ast::{rustname}::{cons.name} (ast::{enumname}{cons.name} {{", depth + 1)
+                if self.typeinfo[cons.name].has_userdata:
+                    generics = "::<LocationRange>"
+                else:
+                    generics = ""
+                self.emit(
+                    f"ast::{rustname}::{cons.name} (ast::{enumname}{cons.name}{generics} {{",
+                    depth + 1,
+                )
                 self.gen_construction_fields(cons, sumname, depth + 1)
                 self.emit("})", depth + 1)
             else:
@@ -691,7 +715,7 @@ def gen_construction(self, cons_path, cons, name, depth):
     def extract_location(self, typename, depth):
         row = self.decode_field(asdl.Field("int", "lineno"), typename)
         column = self.decode_field(asdl.Field("int", "col_offset"), typename)
-        self.emit(f"let _location = ast::Location::new({row}, {column});", depth)
+        self.emit(f"let _location = Location::new({row}, {column});", depth)
 
     def decode_field(self, field, typename):
         name = json.dumps(field.name)
@@ -717,86 +741,20 @@ def write_ast_def(mod, typeinfo, f):
             """
         #![allow(clippy::derive_partial_eq_without_eq)]
         
-        pub use crate::constant::*;
-        pub use rustpython_compiler_core::text_size::{TextSize, TextRange};
+        pub use crate::{Located, constant::*};
+        pub use rustpython_compiler_core::{text_size::{TextSize, TextRange}};
 
         type Ident = String;
         \n
     """
         )
     )
-    StructVisitor(f, typeinfo).emit_attrs(0)
-    f.write(
-        textwrap.dedent(
-            """
-        pub struct Located<T, U = ()> {
-            pub range: TextRange,
-            pub custom: U,
-            pub node: T,
-        }
-    
-        impl<T> Located<T> {
-            pub fn new(start: TextSize, end: TextSize, node: T) -> Self {
-                Self { range: TextRange::new(start, end), custom: (), node }
-            }
-            
-            /// Creates a new node that spans the position specified by `range`.
-            pub fn with_range(node: T, range: TextRange) -> Self {
-                Self {
-                    range,
-                    custom: (),
-                    node,
-                }
-            }
-        
-            /// Returns the absolute start position of the node from the beginning of the document.
-            #[inline]
-            pub const fn start(&self) -> TextSize {
-                self.range.start()
-            }
-        
-            /// Returns the node
-            #[inline]
-            pub fn node(&self) -> &T {
-                &self.node
-            }
-        
-            /// Consumes self and returns the node.
-            #[inline]
-            pub fn into_node(self) -> T {
-                self.node
-            }
-        
-            /// Returns the `range` of the node. The range offsets are absolute to the start of the document.
-            #[inline]
-            pub const fn range(&self) -> TextRange {
-                self.range
-            }
-        
-            /// Returns the absolute position at which the node ends in the source document.
-            #[inline]
-            pub const fn end(&self) -> TextSize {
-                self.range.end()
-            }
-        }
-        
-        impl<T, U> std::ops::Deref for Located<T, U> {
-            type Target = T;
-            
-            fn deref(&self) -> &Self::Target {
-                &self.node
-            }
-        }
-        \n
-    """.lstrip()
-        )
-    )
 
     c = ChainOfVisitors(StructVisitor(f, typeinfo), FoldModuleVisitor(f, typeinfo))
     c.visit(mod)
 
 
-def write_ast_mod(mod, f):
+def write_ast_mod(mod, typeinfo, f):
     f.write(
         textwrap.dedent(
             """
@@ -809,7 +767,11 @@ def write_ast_mod(mod, f):
         )
     )
 
-    c = ChainOfVisitors(ClassDefVisitor(f), TraitImplVisitor(f), ExtendModuleVisitor(f))
+    c = ChainOfVisitors(
+        ClassDefVisitor(f, typeinfo),
+        TraitImplVisitor(f, typeinfo),
+        ExtendModuleVisitor(f, typeinfo),
+    )
     c.visit(mod)
 
 
@@ -830,7 +792,7 @@ def main(input_filename, ast_mod_filename, ast_def_filename, dump_module=False):
         write_ast_def(mod, typeinfo, def_file)
 
         mod_file.write(auto_gen_msg)
-        write_ast_mod(mod, mod_file)
+        write_ast_mod(mod, typeinfo, mod_file)
 
     print(f"{ast_def_filename}, {ast_mod_filename} regenerated.")
 

ast/src/ast_gen.rs

@@ -2,75 +2,11 @@
 
 #![allow(clippy::derive_partial_eq_without_eq)]
 
-pub use crate::constant::*;
+pub use crate::{constant::*, Located};
 pub use rustpython_compiler_core::text_size::{TextRange, TextSize};
 
 type Ident = String;
 
-#[derive(Clone, Debug, PartialEq)]
-pub struct Located<T, U = ()> {
-    pub range: TextRange,
-    pub custom: U,
-    pub node: T,
-}
-
-impl<T> Located<T> {
-    pub fn new(start: TextSize, end: TextSize, node: T) -> Self {
-        Self {
-            range: TextRange::new(start, end),
-            custom: (),
-            node,
-        }
-    }
-
-    /// Creates a new node that spans the position specified by `range`.
-    pub fn with_range(node: T, range: TextRange) -> Self {
-        Self {
-            range,
-            custom: (),
-            node,
-        }
-    }
-
-    /// Returns the absolute start position of the node from the beginning of the document.
-    #[inline]
-    pub const fn start(&self) -> TextSize {
-        self.range.start()
-    }
-
-    /// Returns the node
-    #[inline]
-    pub fn node(&self) -> &T {
-        &self.node
-    }
-
-    /// Consumes self and returns the node.
-    #[inline]
-    pub fn into_node(self) -> T {
-        self.node
-    }
-
-    /// Returns the `range` of the node. The range offsets are absolute to the start of the document.
-    #[inline]
-    pub const fn range(&self) -> TextRange {
-        self.range
-    }
-
-    /// Returns the absolute position at which the node ends in the source document.
-    #[inline]
-    pub const fn end(&self) -> TextSize {
-        self.range.end()
-    }
-}
-
-impl<T, U> std::ops::Deref for Located<T, U> {
-    type Target = T;
-
-    fn deref(&self) -> &Self::Target {
-        &self.node
-    }
-}
-
 #[derive(Clone, Debug, PartialEq)]
 pub struct ModModule<U = ()> {
     pub body: Vec<Stmt<U>>,
@@ -1086,6 +1022,17 @@ pub mod fold {
         type TargetU;
         type Error;
         fn map_user(&mut self, user: U) -> Result<Self::TargetU, Self::Error>;
+        fn map_located<T>(
+            &mut self,
+            located: Located<T, U>,
+        ) -> Result<Located<T, Self::TargetU>, Self::Error> {
+            let custom = self.map_user(located.custom)?;
+            Ok(Located {
+                range: located.range,
+                custom,
+                node: located.node,
+            })
+        }
         fn fold_mod(&mut self, node: Mod<U>) -> Result<Mod<Self::TargetU>, Self::Error> {
             fold_mod(self, node)
         }
@@ -1167,8 +1114,9 @@ pub mod fold {
         node: Located<T, U>,
         f: impl FnOnce(&mut F, T) -> Result<MT, F::Error>,
     ) -> Result<Located<MT, F::TargetU>, F::Error> {
+        let node = folder.map_located(node)?;
         Ok(Located {
-            custom: folder.map_user(node.custom)?,
+            custom: node.custom,
             range: node.range,
             node: f(folder, node.node)?,
         })

ast/src/fold_helpers.rs

@@ -1,6 +1,6 @@
 use crate::{constant, fold::Fold};
 
-pub(crate) trait Foldable<T, U> {
+pub trait Foldable<T, U> {
     type Mapped;
     fn fold<F: Fold<T, TargetU = U> + ?Sized>(
         self,

ast/src/lib.rs

@@ -3,9 +3,16 @@ mod constant;
 #[cfg(feature = "fold")]
 mod fold_helpers;
 mod impls;
+mod located;
+#[cfg(feature = "location")]
+mod locator;
+
 #[cfg(feature = "unparse")]
 mod unparse;
 
 pub use ast_gen::*;
+pub use located::Located;
+#[cfg(feature = "location")]
+pub use locator::Locator;
 
 pub type Suite<U = ()> = Vec<Stmt<U>>;