module Steep
class Source
attr_reader :buffer
attr_reader :path
attr_reader :node
attr_reader :mapping
attr_reader :comments
extend NodeHelper
def initialize(buffer:, path:, node:, mapping:, comments:)
@buffer = buffer
@path = path
@node = node
@mapping = mapping
@comments = comments
end
class Builder < ::Parser::Builders::Default
def string_value(token)
value(token)
end
self.emit_lambda = true
self.emit_procarg0 = true
self.emit_kwargs = true
self.emit_forward_arg = true
end
def self.new_parser
::Parser::Ruby32.new(Builder.new).tap do |parser|
parser.diagnostics.all_errors_are_fatal = true
parser.diagnostics.ignore_warnings = true
end
end
def self.parse(source_code, path:, factory:)
buffer = ::Parser::Source::Buffer.new(path.to_s, 1, source: source_code)
node, comments = new_parser().parse_with_comments(buffer)
# @type var annotations: Array[AST::Annotation::t]
annotations = []
# @type var type_comments: Hash[Integer, type_comment]
type_comments = {}
buffer = RBS::Buffer.new(name: path, content: source_code)
annotation_parser = AnnotationParser.new(factory: factory)
comments = comments.sort_by do |comment|
comment.loc.expression.begin_pos
end
comments.each do |comment|
if comment.inline?
content = comment.text.delete_prefix('#')
content.lstrip!
prefix = comment.text.size - content.size
content.rstrip!
suffix = comment.text.size - content.size - prefix
location = RBS::Location.new(
buffer: buffer,
start_pos: comment.location.expression.begin_pos + prefix,
end_pos: comment.location.expression.end_pos - suffix
)
case
when annotation = annotation_parser.parse(content, location: location)
annotations << annotation
when assertion = AST::Node::TypeAssertion.parse(location)
type_comments[assertion.line] = assertion
when tapp = AST::Node::TypeApplication.parse(location)
type_comments[tapp.line] = tapp
end
end
end
map = {} #: Hash[Parser::AST::Node, Array[AST::Annotation::t]]
map.compare_by_identity
if node
node = insert_type_node(node, type_comments)
construct_mapping(node: node, annotations: annotations, mapping: map)
end
annotations.each do |annot|
map[node] ||= []
map[node] << annot
end
new(buffer: buffer, path: path, node: node, mapping: map, comments: comments)
end
def self.construct_mapping(node:, annotations:, mapping:, line_range: nil)
case node.type
when :if
cond_node, truthy_node, falsy_node, loc = deconstruct_if_node!(node)
if node.loc.respond_to?(:question)
# Skip ternary operator
each_child_node node do |child|
construct_mapping(node: child, annotations: annotations, mapping: mapping, line_range: nil)
end
else
if_loc = loc #: NodeHelper::condition_loc
if if_loc.expression.begin_pos == if_loc.keyword.begin_pos
construct_mapping(node: cond_node,annotations: annotations, mapping: mapping, line_range: nil)
if truthy_node
if if_loc.keyword.source == "if" || if_loc.keyword.source == "elsif"
# if foo
# bar <=
# end
then_start = if_loc.begin&.last_line || cond_node.loc.last_line
then_end = if_loc.else&.line || if_loc.last_line
else
# unless foo
# else
# bar <=
# end
if_loc.else or raise
then_start = if_loc.else.last_line
then_end = loc.last_line
end
construct_mapping(node: truthy_node, annotations: annotations, mapping: mapping, line_range: then_start...then_end)
end
if falsy_node
if if_loc.keyword.source == "if" || if_loc.keyword.source == "elsif"
# if foo
# else
# bar <=
# end
if_loc.else or raise
else_start = if_loc.else.last_line
else_end = if_loc.last_line
else
# unless foo
# bar <=
# end
else_start = if_loc.begin&.last_line || cond_node.loc.last_line
else_end = if_loc.else&.line || if_loc.last_line
end
construct_mapping(node: falsy_node, annotations: annotations, mapping: mapping, line_range: else_start...else_end)
end
else
# postfix if/unless
each_child_node(node) do |child|
construct_mapping(node: child, annotations: annotations, mapping: mapping, line_range: nil)
end
end
end
when :while, :until
cond_node, body_node, loc = deconstruct_whileish_node!(node)
if loc.expression.begin_pos == loc.keyword.begin_pos
# prefix while
loc.end or raise
construct_mapping(node: cond_node, annotations: annotations, mapping: mapping, line_range: nil)
if body_node
body_start = cond_node.loc.last_line
body_end = loc.end.line
construct_mapping(node: body_node, annotations: annotations, mapping: mapping, line_range: body_start...body_end)
end
else
# postfix while
each_child_node(node) do |child|
construct_mapping(node: child, annotations: annotations, mapping: mapping, line_range: nil)
end
end
when :while_post, :until_post
cond_node, body_node, loc = deconstruct_whileish_node!(node)
construct_mapping(node: cond_node, annotations: annotations, mapping: mapping, line_range: nil)
if body_node
body_start = loc.expression.line
body_end = loc.keyword.line
construct_mapping(node: body_node, annotations: annotations, mapping: mapping, line_range: body_start...body_end)
end
when :case
cond_node, when_nodes, else_node, loc = deconstruct_case_node!(node)
if cond_node
construct_mapping(node: cond_node, annotations: annotations, mapping: mapping, line_range: nil)
end
if else_node
loc.else or raise
loc.end or raise
else_start = loc.else.last_line
else_end = loc.end.line
construct_mapping(node: else_node, annotations: annotations, mapping: mapping, line_range: else_start...else_end)
end
when_nodes.each do |child|
if child.type == :when
construct_mapping(node: child, annotations: annotations, mapping: mapping, line_range: nil)
end
end
when :when
operands, body, loc = deconstruct_when_node!(node)
last_cond = operands.last or raise
operands.each do |child|
construct_mapping(node: child, annotations: annotations, mapping: mapping, line_range: nil)
end
if body
cond_end = loc.begin&.last_line || last_cond.loc.last_line+1
body_end = body.loc.last_line
construct_mapping(node: body, annotations: annotations, mapping: mapping, line_range: cond_end...body_end)
end
when :rescue
body, resbodies, else_node, loc = deconstruct_rescue_node!(node)
if else_node
loc.else or raise
else_start = loc.else.last_line
else_end = loc.last_line
construct_mapping(node: else_node, annotations: annotations, mapping: mapping, line_range: else_start...else_end)
end
each_child_node(node) do |child|
construct_mapping(node: child, annotations: annotations, mapping: mapping, line_range: nil)
end
else
each_child_node(node) do |child|
construct_mapping(node: child, annotations: annotations, mapping: mapping, line_range: nil)
end
end
associated_annotations, other_annotations = annotations.partition do |annot|
location = node.loc
annot.line or next
case node.type
when :def, :module, :class, :block, :ensure, :defs, :resbody
location = node.loc
location.line <= annot.line && annot.line < location.last_line
else
if line_range
line_range.begin <= annot.line && annot.line < line_range.end
end
end
end
associated_annotations.each do |annot|
mapping[node] ||= []
mapping[node] << annot
end
annotations.replace(other_annotations)
end
def self.map_child_node(node, type = nil, skip: nil)
children = node.children.map do |child|
if child.is_a?(Parser::AST::Node)
if skip
if skip.member?(child)
child
else
yield child
end
else
yield child
end
else
child
end
end
node.updated(type, children)
end
def annotations(block:, factory:, context:)
AST::Annotation::Collection.new(
annotations: (mapping[block] || []),
factory: factory,
context: context
)
end
def each_annotation(&block)
if block_given?
mapping.each do |node, annots|
yield [node, annots]
end
else
enum_for :each_annotation
end
end
def each_heredoc_node(node = self.node, parents = [], &block)
if block
return unless node
case node.type
when :dstr, :str
if node.location.respond_to?(:heredoc_body)
yield [[node, *parents], _ = node.location]
end
end
parents.unshift(node)
Source.each_child_node(node) do |child|
each_heredoc_node(child, parents, &block)
end
parents.shift()
else
enum_for :each_heredoc_node, node
end
end
def find_heredoc_nodes(line, column, position)
each_heredoc_node() do |nodes, location|
node = nodes[0]
loc = location.heredoc_body #: Parser::Source::Range
if range = loc.to_range
if range.begin <= position && position <= range.end
return nodes
end
end
end
nil
end
def find_nodes_loc(node, position, parents)
range = node.location.expression&.to_range
if range
if range.begin <= position && position <= range.end
parents.unshift node
Source.each_child_node(node) do |child|
if ns = find_nodes_loc(child, position, parents)
return ns
end
end
parents
end
end
end
def find_nodes(line:, column:)
return [] unless node
position = buffer.loc_to_pos([line, column])
if heredoc_nodes = find_heredoc_nodes(line, column, position)
Source.each_child_node(heredoc_nodes[0]) do |child|
if nodes = find_nodes_loc(child, position, heredoc_nodes)
return nodes
end
end
return heredoc_nodes
else
find_nodes_loc(node, position, [])
end
end
def find_comment(line:, column:)
pos = buffer.loc_to_pos([line, column])
comment = comments.bsearch do |comment|
pos <= comment.loc.expression.end_pos
end
if comment
if comment.loc.expression.begin_pos < pos
comment
end
end
end
def self.delete_defs(node, allow_list)
case node.type
when :def
if allow_list.include?(node)
node
else
node.updated(:nil, [])
end
when :defs
if allow_list.include?(node)
node
else
delete_defs(node.children[0], allow_list)
end
else
map_child_node(node) do |child|
delete_defs(child, allow_list)
end
end
end
def without_unrelated_defs(line:, column:)
if node
nodes = find_nodes(line: line, column: column) || []
defs = Set[].compare_by_identity.merge(nodes.select {|node| node.type == :def || node.type == :defs })
node_ = Source.delete_defs(node, defs)
# @type var mapping: Hash[Parser::AST::Node, Array[AST::Annotation::t]]
mapping = {}
mapping.compare_by_identity
annotations = self.mapping.values.flatten
Source.construct_mapping(node: node_, annotations: annotations, mapping: mapping)
annotations.each do |annot|
mapping[node_] ||= []
mapping[node_] << annot
end
Source.new(buffer: buffer, path: path, node: node_, mapping: mapping, comments: comments)
else
self
end
end
def self.insert_type_node(node, comments)
if node.location.expression
first_line = node.location.expression.first_line
last_line = node.location.expression.last_line
last_comment = comments[last_line]
if (first_line..last_line).none? {|l| comments.key?(l) }
return node
end
case
when last_comment.is_a?(AST::Node::TypeAssertion)
case node.type
when :lvasgn, :ivasgn, :gvasgn, :cvasgn, :casgn
# Skip
when :masgn
lhs, rhs = node.children
node = node.updated(nil, [lhs, insert_type_node(rhs, comments)])
return adjust_location(node)
when :return, :break, :next
# Skip
when :begin
location = node.loc #: Parser::Source::Map & Parser::AST::_Collection
if location.begin
# paren
child_assertions = comments.except(last_line)
node = map_child_node(node) {|child| insert_type_node(child, child_assertions) }
node = adjust_location(node)
return assertion_node(node, last_comment)
end
else
child_assertions = comments.except(last_line)
node = map_child_node(node) {|child| insert_type_node(child, child_assertions) }
node = adjust_location(node)
return assertion_node(node, last_comment)
end
when selector_line = sendish_node?(node)
if (comment = comments[selector_line]).is_a?(AST::Node::TypeApplication)
child_assertions = comments.except(selector_line)
case node.type
when :block
send, *children = node.children
node = node.updated(
nil,
[
map_child_node(send) {|child| insert_type_node(child, child_assertions) },
*children.map do |child|
if child
insert_type_node(child, child_assertions)
end
end
]
)
when :numblock
send, size, body = node.children
node = node.updated(
nil,
[
map_child_node(send) {|child| insert_type_node(child, child_assertions) },
size,
insert_type_node(body, child_assertions)
]
)
else
node = map_child_node(node) {|child| insert_type_node(child, child_assertions) }
end
node = adjust_location(node)
return type_application_node(node, comment)
end
end
end
case node.type
when :class
class_name, super_class, class_body = node.children
adjust_location(
node.updated(
nil,
[
class_name,
super_class,
class_body ? insert_type_node(class_body, comments) : nil
]
)
)
when :module
module_name, module_body = node.children
adjust_location(
node.updated(
nil,
[
module_name,
module_body ? insert_type_node(module_body, comments) : nil
]
)
)
else
adjust_location(
map_child_node(node, nil) {|child| insert_type_node(child, comments) }
)
end
end
def self.sendish_node?(node)
send_node =
case node.type
when :send, :csend
node
when :block, :numblock
send = node.children[0]
case send.type
when :send, :csend
send
end
end
if send_node
receiver_node, name, _, location = deconstruct_send_node!(send_node)
if receiver_node
if location.dot
location.selector.line
end
else
location.selector.line
end
end
end
def self.adjust_location(node)
if end_pos = node.location.expression&.end_pos
if last_pos = each_child_node(node).map {|node| node.location.expression&.end_pos }.compact.max
if last_pos > end_pos
props = { location: node.location.with_expression(node.location.expression.with(end_pos: last_pos)) }
end
end
end
if props
node.updated(nil, nil, props)
else
node
end
end
def self.assertion_node(node, type)
map = Parser::Source::Map.new(node.location.expression.with(end_pos: type.location.end_pos))
Parser::AST::Node.new(:assertion, [node, type], { location: map })
end
def self.type_application_node(node, tapp)
if node.location.expression.end_pos > tapp.location.end_pos
map = Parser::Source::Map.new(node.location.expression)
else
map = Parser::Source::Map.new(node.location.expression.with(end_pos: tapp.location.end_pos))
end
node = Parser::AST::Node.new(:tapp, [node, tapp], { location: map })
tapp.set_node(node)
node
end
end
end