class Steep::Typing

def add_context_for_body(node, context:)
  case node.type
  when :class
    name_node, super_node, _ = node.children
    begin_pos = if super_node
                  super_node.loc.expression.end_pos
                else
                  name_node.loc.expression.end_pos
                end
    end_pos = node.loc.end.begin_pos
    add_context(begin_pos..end_pos, context: context)
  when :module
    name_node = node.children[0]
    begin_pos = name_node.loc.expression.end_pos
    end_pos = node.loc.end.begin_pos
    add_context(begin_pos..end_pos, context: context)
  when :def, :defs
    args_node = case node.type
                when :def
                  node.children[1]
                when :defs
                  node.children[2]
                end
    body_begin_pos = if args_node.loc.expression
                       args_node.loc.expression.end_pos
                     else
                       node.loc.name.end_pos
                     end
    body_end_pos = node.loc.end.begin_pos
    add_context(body_begin_pos..body_end_pos, context: context)
  when :block
    send_node, args_node, _ = node.children
    begin_pos = if send_node.type != :lambda && args_node.loc.expression
                  args_node.loc.expression.end_pos
                else
                  node.loc.begin.end_pos
                end
    end_pos = node.loc.end.begin_pos
    add_context(begin_pos..end_pos, context: context)
  else
    raise "Unexpected node for insert_context: #{node.type}"
  end
end