Vernacular::AST

gem 'vernacular-ast'

Vernacular::ASTModifier.new do |modifier|
  # Extend the parser to support and equal sign and a class path following the
  # declaration of a functions arguments to represent its return type.
  modifier.extend_parser(:f_arglist, 'f_arglist tEQL cpath', <<~PARSE)
    result = @builder.type_check_arglist(*val)
  PARSE

  # Extend the builder by adding a `type_check_arglist` function that will build
  # a new node type and place it at the end of the argument list.
  modifier.extend_builder(:type_check_arglist) do |arglist, equal, cpath|
    arglist << n(:type_check_arglist, [equal, cpath], nil)
  end

  # Extend the rewriter by adding an `on_def` callback, which will be called
  # whenever a `def` node is added to the AST. Then, loop through and find any
  # `type_check_arglist` nodes, and remove them. Finally, insert the
  # appropriate raises around the execution of the function to mirror the type
  # checking.
  modifier.build_rewriter do
    def on_def(node)
      type_check_node = node.children[1].children.last
      return super if !type_check_node || type_check_node.type != :type_check_arglist

      remove(type_check_node.children[0][1])
      remove(type_check_node.children[1].loc.expression)
      type = build_constant(type_check_node.children[1])

      @source_rewriter.transaction do
        insert_before(node.children[2].loc.expression, "result = begin\n")
        insert_after(node.children[2].loc.expression,
          "\nend\nraise \"Invalid return value, expected #{type}, " <<
          "got \#{result.class.name}\" unless result.is_a?(#{type})\nresult")
      end

      super
    end

    private

    def build_constant(node, suffix = nil)
      child_node, name = node.children
      new_name = suffix ? "#{name}::#{suffix}" : name
      child_node ? build_constant(child_node, new_name) : new_name
    end
  end
end