Merge pull request #3044 from ruby/owned-ast-locations

Add the missing source locations to the owned AST

Author: soutaro

rust/ruby-rbs/src/ast/convert.rs

@@ -12,10 +12,11 @@ use crate::ast::location::{
     AliasDeclarationLocation, AliasLocation, AliasMemberLocation, AttributeMemberLocation,
     ClassDeclarationLocation, ClassInstanceLocation, ClassSingletonLocation, ClassSuperLocation,
     ConstantDeclarationLocation, FunctionParamLocation, GlobalDeclarationLocation,
-    InterfaceDeclarationLocation, InterfaceLocation, LocationRange, MethodDefinitionLocation,
-    MethodTypeLocation, MixinMemberLocation, ModuleDeclarationLocation, ModuleSelfLocation,
-    TypeAliasDeclarationLocation, TypeParamLocation, UseDirectiveLocation, UseSingleClauseLocation,
-    UseWildcardClauseLocation, VariableMemberLocation,
+    InterfaceDeclarationLocation, InterfaceLocation, KeywordParamLocation, LocationRange,
+    MethodDefinitionLocation, MethodTypeLocation, MixinMemberLocation, ModuleDeclarationLocation,
+    ModuleSelfLocation, RecordFieldLocation, TypeAliasDeclarationLocation, TypeParamLocation,
+    UseDirectiveLocation, UseSingleClauseLocation, UseWildcardClauseLocation,
+    VariableMemberLocation,
 };
 use crate::ast::members::{
     AliasKind, AliasMember, AttrAccessorMember, AttrReaderMember, AttrWriterMember, AttributeKind,
@@ -163,6 +164,7 @@ impl<'a> AstConverter<'a> {
             Node::RecordType(node) => {
                 let mut fields = Vec::new();
                 for (key, value) in node.all_fields().iter() {
+                    let key_range = key.location();
                     let key = self.convert_record_key(&key);
                     let Node::RecordFieldType(field) = value else {
                         panic_expected("record field value while converting record type", &value);
@@ -171,6 +173,7 @@ impl<'a> AstConverter<'a> {
                         key,
                         ty: self.convert_type(&field.type_()),
                         required: field.required(),
+                        location: Some(record_field_location(key_range, field.location())),
                     });
                 }
                 Type::Record(RecordType {
@@ -466,6 +469,7 @@ impl<'a> AstConverter<'a> {
         MethodDefinitionOverload {
             method_type: self.convert_method_type_node(&node.method_type()),
             annotations: self.convert_annotations(node.annotations()),
+            location: Some(convert_range(node.location())),
         }
     }
 
@@ -929,6 +933,7 @@ impl<'a> AstConverter<'a> {
                 KeywordParam {
                     name: self.intern_symbol(&symbol),
                     param: self.convert_function_param_node(&value),
+                    location: Some(keyword_param_location(symbol.location(), value.location())),
                 }
             })
             .collect()
@@ -1068,6 +1073,44 @@ fn convert_optional_range(range: Option<RBSLocationRange>) -> Option<LocationRan
     range.map(convert_range)
 }
 
+/// Builds the location of a keyword parameter, which the C AST keeps as a pair
+/// of a key symbol and a function param node instead of a single node.
+fn keyword_param_location(
+    name_range: RBSLocationRange,
+    param_range: RBSLocationRange,
+) -> KeywordParamLocation {
+    let name_range = convert_range(name_range);
+
+    KeywordParamLocation {
+        range: span_range(name_range, convert_range(param_range)),
+        name_range,
+    }
+}
+
+/// Builds the location of a record field, which the C AST keeps as a pair of a
+/// key node and a field type node instead of a single node.
+fn record_field_location(
+    key_range: RBSLocationRange,
+    field_range: RBSLocationRange,
+) -> RecordFieldLocation {
+    let key_range = convert_range(key_range);
+
+    RecordFieldLocation {
+        range: span_range(key_range, convert_range(field_range)),
+        key_range,
+    }
+}
+
+/// Returns the range starting where `start` starts and ending where `end` ends.
+fn span_range(start: LocationRange, end: LocationRange) -> LocationRange {
+    LocationRange {
+        start_char: start.start_char,
+        start_byte: start.start_byte,
+        end_char: end.end_char,
+        end_byte: end.end_byte,
+    }
+}
+
 fn variable_member_location(
     range: RBSLocationRange,
     name_range: RBSLocationRange,

rust/ruby-rbs/src/ast/location.rs

@@ -70,6 +70,40 @@ pub struct FunctionParamLocation {
     pub name_range: Option<LocationRange>,
 }
 
+/// ```rbs
+/// (name: String) -> void
+///  ^^^^^^^^^^^^            range
+///  ^^^^                    name
+///
+/// (?size: Integer bytes) -> void
+///   ^^^^^^^^^^^^^^^^^^^            range
+///   ^^^^                           name
+/// ```
+///
+/// The `?` marker of an optional keyword is not part of `range`: it belongs to
+/// the enclosing function type, not to the keyword parameter itself.
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
+pub struct KeywordParamLocation {
+    pub range: LocationRange,
+    pub name_range: LocationRange,
+}
+
+/// ```rbs
+/// { name: String, "id" => Integer }
+///   ^^^^^^^^^^^^                     range
+///   ^^^^                             key
+///                 ^^^^^^^^^^^^^^^    range
+///                 ^^^^               key
+/// ```
+///
+/// The `?` marker of an optional field is not part of `range`: it belongs to
+/// the enclosing record type, not to the field itself.
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
+pub struct RecordFieldLocation {
+    pub range: LocationRange,
+    pub key_range: LocationRange,
+}
+
 /// ```rbs
 /// _Foo
 /// ^^^^ name

rust/ruby-rbs/src/ast/members.rs

@@ -1,8 +1,8 @@
 use crate::ast::annotation::Annotation;
 use crate::ast::comment::Comment;
 use crate::ast::location::{
-    AliasMemberLocation, AttributeMemberLocation, MethodDefinitionLocation, MixinMemberLocation,
-    VariableMemberLocation,
+    AliasMemberLocation, AttributeMemberLocation, LocationRange, MethodDefinitionLocation,
+    MixinMemberLocation, VariableMemberLocation,
 };
 use crate::ast::method_type::MethodType;
 use crate::ast::types::Type;
@@ -54,6 +54,9 @@ pub struct MethodDefinitionMember {
 pub struct MethodDefinitionOverload {
     pub method_type: MethodType,
     pub annotations: Vec<Annotation>,
+    /// Starts at the `:` or `|` separator preceding the overload, so that the
+    /// overloads of a method definition tile its whole type without gaps.
+    pub location: Option<LocationRange>,
 }
 
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]

rust/ruby-rbs/src/ast/mod.rs

@@ -34,11 +34,11 @@ pub use location::{
     AliasDeclarationLocation, AliasLocation, AliasMemberLocation, AttributeMemberLocation,
     ClassDeclarationLocation, ClassInstanceLocation, ClassSingletonLocation, ClassSuperLocation,
     ConstantDeclarationLocation, FunctionParamLocation, GlobalDeclarationLocation,
-    InterfaceDeclarationLocation, InterfaceLocation, LocationRange, MethodDefinitionLocation,
-    MethodTypeLocation, MixinMemberLocation, ModuleDeclarationLocation, ModuleSelfLocation,
-    ResolveTypeNamesDirectiveLocation, TypeAliasDeclarationLocation, TypeParamLocation,
-    UseDirectiveLocation, UseSingleClauseLocation, UseWildcardClauseLocation,
-    VariableMemberLocation,
+    InterfaceDeclarationLocation, InterfaceLocation, KeywordParamLocation, LocationRange,
+    MethodDefinitionLocation, MethodTypeLocation, MixinMemberLocation, ModuleDeclarationLocation,
+    ModuleSelfLocation, RecordFieldLocation, ResolveTypeNamesDirectiveLocation,
+    TypeAliasDeclarationLocation, TypeParamLocation, UseDirectiveLocation, UseSingleClauseLocation,
+    UseWildcardClauseLocation, VariableMemberLocation,
 };
 pub use members::{
     AliasKind, AliasMember, AttrAccessorMember, AttrReaderMember, AttrWriterMember, AttributeKind,
@@ -58,8 +58,9 @@ pub use types::{
 #[cfg(test)]
 mod tests {
     use crate::ast::{
-        AstConverter, BaseType, BaseTypeKind, ClassMember, Declaration, Directive, IvarName,
-        Literal, Member, MethodKind, ModuleMember, RecordKey, Type, UseClause,
+        AstConverter, BaseType, BaseTypeKind, ClassMember, Declaration, Directive, Function,
+        IvarName, Literal, LocationRange, Member, MethodKind, ModuleMember, RecordKey, Type,
+        UseClause,
     };
     use crate::interner::StringInterner;
     use crate::node::{Node, parse};
@@ -324,4 +325,101 @@ mod tests {
         assert_eq!(type_names.display(wildcard.namespace, &strings), "Foo::Baz");
         assert!(wildcard.location.is_some());
     }
+
+    #[test]
+    fn converts_keyword_param_locations() {
+        let source = "class Foo\n  def bar: (name: String, ?size: Integer bytes) -> void\nend\n";
+        let signature = parse(source).unwrap();
+
+        let mut strings = StringInterner::new();
+        let mut type_names = TypeNameInterner::new();
+        let mut converter = AstConverter::new(&mut strings, &mut type_names);
+        let declaration =
+            converter.convert_declaration(&signature.declarations().iter().next().unwrap());
+
+        let Declaration::Class(class_decl) = &declaration else {
+            panic!("expected class declaration");
+        };
+        let ClassMember::Member(Member::MethodDefinition(method)) = &class_decl.members[0] else {
+            panic!("expected method definition member");
+        };
+        let Function::Typed(function) = &method.overloads[0].method_type.function else {
+            panic!("expected typed function");
+        };
+
+        let text =
+            |range: &LocationRange| &source[range.start_byte as usize..range.end_byte as usize];
+
+        let required = &function.required_keywords[0];
+        let required_location = required.location.as_ref().unwrap();
+        assert_eq!(required.name, strings.intern("name"));
+        assert_eq!(text(&required_location.range), "name: String");
+        assert_eq!(text(&required_location.name_range), "name");
+
+        let optional = &function.optional_keywords[0];
+        let optional_location = optional.location.as_ref().unwrap();
+        assert_eq!(optional.name, strings.intern("size"));
+        assert_eq!(text(&optional_location.range), "size: Integer bytes");
+        assert_eq!(text(&optional_location.name_range), "size");
+    }
+
+    #[test]
+    fn converts_record_field_locations() {
+        let source = "type t = { name: String, ?age: Integer, \"id\" => Integer }\n";
+        let signature = parse(source).unwrap();
+
+        let mut strings = StringInterner::new();
+        let mut type_names = TypeNameInterner::new();
+        let mut converter = AstConverter::new(&mut strings, &mut type_names);
+        let declaration =
+            converter.convert_declaration(&signature.declarations().iter().next().unwrap());
+
+        let Declaration::TypeAlias(alias) = &declaration else {
+            panic!("expected type alias declaration");
+        };
+        let Type::Record(record) = &alias.ty else {
+            panic!("expected record type");
+        };
+
+        let text =
+            |range: &LocationRange| &source[range.start_byte as usize..range.end_byte as usize];
+        let location = |index: usize| record.fields[index].location.as_ref().unwrap();
+
+        assert_eq!(text(&location(0).range), "name: String");
+        assert_eq!(text(&location(0).key_range), "name");
+
+        // The `?` marker belongs to the record type, not to the field.
+        assert_eq!(text(&location(1).range), "age: Integer");
+        assert_eq!(text(&location(1).key_range), "age");
+
+        assert_eq!(text(&location(2).range), "\"id\" => Integer");
+        assert_eq!(text(&location(2).key_range), "\"id\"");
+    }
+
+    #[test]
+    fn converts_method_definition_overload_locations() {
+        let source = "class Foo\n  def bar: () -> void\n        | (Integer) -> String\nend\n";
+        let signature = parse(source).unwrap();
+
+        let mut strings = StringInterner::new();
+        let mut type_names = TypeNameInterner::new();
+        let mut converter = AstConverter::new(&mut strings, &mut type_names);
+        let declaration =
+            converter.convert_declaration(&signature.declarations().iter().next().unwrap());
+
+        let Declaration::Class(class_decl) = &declaration else {
+            panic!("expected class declaration");
+        };
+        let ClassMember::Member(Member::MethodDefinition(method)) = &class_decl.members[0] else {
+            panic!("expected method definition member");
+        };
+
+        let text =
+            |range: &LocationRange| &source[range.start_byte as usize..range.end_byte as usize];
+        let location = |index: usize| method.overloads[index].location.as_ref().unwrap();
+
+        // The leading `:`/`|` separator is part of the overload range.
+        assert_eq!(text(location(0)), ": () -> void");
+        assert_eq!(text(location(1)), "| (Integer) -> String");
+    }
 }

rust/ruby-rbs/src/ast/types.rs

@@ -1,6 +1,6 @@
 use crate::ast::location::{
     AliasLocation, ClassInstanceLocation, ClassSingletonLocation, FunctionParamLocation,
-    InterfaceLocation, LocationRange,
+    InterfaceLocation, KeywordParamLocation, LocationRange, RecordFieldLocation,
 };
 use crate::ids::{SymbolId, TypeName};
 
@@ -99,6 +99,7 @@ pub struct RecordField {
     pub key: RecordKey,
     pub ty: Type,
     pub required: bool,
+    pub location: Option<RecordFieldLocation>,
 }
 
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
@@ -136,6 +137,7 @@ pub struct FunctionType {
 pub struct KeywordParam {
     pub name: SymbolId,
     pub param: FunctionParam,
+    pub location: Option<KeywordParamLocation>,
 }
 
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]