LaTeXify Ruby methods

latexify.rb

#!/usr/bin/env ruby
# frozen_string_literal: true

require "bundler/inline"

gemfile do
  source "https://rubygems.org"
  gem "syntax_tree"
end

# These are names of variables associated with math constants that are
# understood by LaTeX. We replace them with the corresponding LaTeX command.
MATH_SYMBOLS = [
  "aleph", "alpha", "beta", "beth", "chi", "daleth", "delta", "digamma",
  "epsilon", "eta", "gamma", "gimel", "iota", "kappa", "lambda", "mu", "nu",
  "omega", "omega", "phi", "pi", "psi", "rho", "sigma", "tau", "theta",
  "upsilon", "varepsilon", "varkappa", "varphi", "varpi", "varrho", "varsigma",
  "vartheta", "xi", "zeta", "Delta", "Gamma", "Lambda", "Omega", "Phi", "Pi",
  "Sigma", "Theta", "Upsilon", "Xi"
]

# These are methods that live on the Math module that have corresponding
# representations in LaTeX. The keys are the names of the methods and the values
# are pairs that should be placed around the expression being passed as the only
# argument.
MATH_METHODS = {
  acos: ["\\arccos{\\left({", "}\\right)}"],
  acosh: ["\\mathrm{arccosh}{\\left({", "}\\right)}"],
  asin: ["\\arcsin{\\left({", "}\\right)}"],
  asinh: ["\\mathrm{arcsinh}{\\left({", "}\\right)}"],
  atan: ["\\arctan{\\left({", "}\\right)}"],
  atanh: ["\\mathrm{arctanh}{\\left({", "}\\right)}"],
  cos: ["\\cos{\\left({", "}\\right)}"],
  cosh: ["\\cosh{\\left({", "}\\right)}"],
  exp: ["\\exp{\\left({", "}\\right)}"],
  gamma: ["\\Gamma\\left({", "}\\right)"],
  log: ["\\log{\\left({", "}\\right)}"],
  log10: ["\\log_{10}{\\left({", "}\\right)}"],
  log2: ["\\log_{2}{\\left({", "}\\right)}"],
  sin: ["\\sin{\\left({", "}\\right)}"],
  sinh: ["\\sinh{\\left({", "}\\right)}"],
  sqrt: ["\\sqrt{", "}"],
  tan: ["\\tan{\\left({", "}\\right)}"],
  tanh: ["\\tanh{\\left({", "}\\right)}"]
}

# These are methods that live on Numeric that have corresponding representations
# in LaTeX. The keys are the names of the methods and the values are pairs that
# should be placed around the receiver.
RECEIVER_METHODS = {
  abs: ["\\left|", "\\right|"],
  ceil: ["\\left\\lceil{", "}\\right\\rceil"],
  floor: ["\\left\\lfloor{", "}\\right\\rfloor"]
}

# This method is responsible for finding the syntax tree that corresponds to the
# given method object. It does this in a couple of steps.
#
# 1. The first is to access the method's source location. This is a tuple of the
#    file path and line number where the method was defined.
# 2. The second is to lex the file at the given path using Ripper. This gives us
#    all of the tokens on that line. We look for the tokens that correspond to
#    the method definition so that we know exactly the column offset of the def
#    keyword. This part isn't really necessary, but it's nice to do for
#    accuracy.
# 3. The third is to parse the file using Syntax Tree and recursively descend
#    the syntax tree until we find the method definition. We do this by using a
#    binary search over the given child nodes.
# 4. Finally, we return the tree if one was found, otherwise we return nil.
def find_method(method)
  filepath, lineno = method.source_location
  name = method.name.to_s

  source = SyntaxTree.read(filepath)

  Ripper.lex(source.lines[lineno - 1]) => [
    *,
    [[_, column], :on_kw, "def", _],
    [[_, _], :on_sp, _, _],
    [[_, _], :on_ident, ^name, _],
    *
  ]

  position = source.lines[0...(lineno - 1)].sum(&:length) + column
  node = SyntaxTree.parse(source)

  while node
    return node if node in SyntaxTree::Def[
      name: SyntaxTree::Ident[value: ^name],
      location: { start_char: ^position }
    ]

    node =
      node.child_nodes.bsearch do |child|
        location = child.location
        if (location.start_char...location.end_char).cover?(position)
          0
        else
          position <=> location.start_char
        end
      end
  end
end

# This method is responsible for converting a node in the syntax tree to an
# equivalent LaTeX expression. For the most part it attempts to match the
# semantics used by the google/latexify_py package here:
#
#     https://github.com/google/latexify_py
#
# There are a couple of differences because of Ruby syntax and also where the
# method definitions for various methods lives. For example in python, abs is a
# globally-accessible function, whereas in Ruby it's a method that lives on
# Numeric.
def latexify(node)
  case node
  in SyntaxTree::Binary[left:, operator: :*, right:]
    while right
      case right
      in SyntaxTree::VarRef[value: SyntaxTree::Ident]
        return "#{latexify(left)}#{latexify(right)}"
      in SyntaxTree::Binary[left: _, operator: :*, right: _]
      else
        return "#{latexify(left)} * #{latexify(right)}"
      end
    end
  in SyntaxTree::Binary[left:, operator: :/, right:]
    left = (left in SyntaxTree::Paren[contents: SyntaxTree::Statements[body: [statement]]]) ? statement : left
    right = (right in SyntaxTree::Paren[contents: SyntaxTree::Statements[body: [statement]]]) ? statement : right
    "\\frac{#{latexify(left)}}{#{latexify(right)}}"
  in SyntaxTree::Binary[left:, operator: :**, right:]
    "#{latexify(left)}^{#{latexify(right)}}"
  in SyntaxTree::Binary[left:, operator:, right:]
    display = {
      :>= => "\\ge",
      :<= => "\\le",
      :!= => "\\ne",
      :== => "=",
      :=== => "\\equiv"
    }.fetch(operator, operator.to_s)

    "#{latexify(left)} #{display} #{latexify(right)}"
  in SyntaxTree::Call[
       receiver: SyntaxTree::VarRef[value: SyntaxTree::Const[value: "Math"]],
       operator: SyntaxTree::Period[value: "."],
       message: SyntaxTree::Ident[value:],
       arguments: SyntaxTree::ArgParen[arguments: SyntaxTree::Args[parts: [argument]]]
     ] if MATH_METHODS.key?(value.to_sym)
    MATH_METHODS[value.to_sym].join(latexify(argument))
  in SyntaxTree::Call[
       receiver:,
       operator: SyntaxTree::Period[value: "."],
       message: SyntaxTree::Ident[value:],
       arguments: nil,
     ] if RECEIVER_METHODS.key?(value.to_sym)
    RECEIVER_METHODS[value.to_sym].join(latexify(receiver))
  in SyntaxTree::Const | SyntaxTree::Ident
    if MATH_SYMBOLS.include?(node.value)
      "{\\#{node.value}}"
    else
      node.value
    end
  in SyntaxTree::Def[
       name: SyntaxTree::Ident[value:],
       params: SyntaxTree::Params => params,
       bodystmt: SyntaxTree::BodyStmt[statements:, rescue_clause: nil, ensure_clause: nil, else_clause: nil]
     ] if params.empty?
    "\\mathrm{#{value}}()\\triangleq #{latexify(statements)}"
  in SyntaxTree::Def[
       name: SyntaxTree::Ident[value:],
       params: SyntaxTree::Paren[contents: SyntaxTree::Params[requireds:, optionals: [], rest: nil, posts: [], keywords: [], keyword_rest: nil, block: nil]],
       bodystmt: SyntaxTree::BodyStmt[statements:, rescue_clause: nil, ensure_clause: nil, else_clause: nil]
     ]
    params = requireds.map { |required| latexify(required) }.join(", ")
    "\\mathrm{#{value}}(#{params})\\triangleq #{latexify(statements)}"
  in SyntaxTree::FCall[value: SyntaxTree::Ident[value:], arguments: SyntaxTree::ArgParen[arguments: SyntaxTree::Args[parts:]]]
    arguments = parts.map { |part| latexify(part) }.join(", ")
    "\\mathrm{#{value}}\\left(#{arguments}\\right)"
  in SyntaxTree::FloatLiteral[value:]
    value
  in SyntaxTree::If[predicate:, statements:, consequent:]
    clauses = [[statements, predicate]]

    while consequent
      case consequent
      in SyntaxTree::Elsif[predicate:, statements:, consequent:]
        clauses << [statements, predicate]
      in SyntaxTree::Else[statements:]
        clauses << [statements, nil]
        consequent = nil
      end
    end

    clauses.map! do |(statements, predicate)|
      display = predicate ? "\\mathrm{if} \\ #{latexify(predicate)}" : "\\mathrm{otherwise}"
      "#{latexify(statements)}, & #{display}"
    end

    "\\left\\{ \\begin{array}{ll} #{clauses.join(" \\\\ ")} \\end{array} \\right."
  in SyntaxTree::Int[value:]
    value
  in SyntaxTree::Not[statement:]
    "\\lnot\\left(#{latexify(statement)}\\right)"
  in SyntaxTree::Paren[contents: SyntaxTree::Statements[body: [statement]]]
    "\\left(#{latexify(statement)}\\right)"
  in SyntaxTree::Return[arguments: SyntaxTree::Args[parts: [part]]]
    latexify(part)
  in SyntaxTree::Statements[body:]
    body.grep_v(SyntaxTree::Comment) => [statement]
    latexify(statement)
  in SyntaxTree::Unary[operator:, statement:]
    left, right = {
      "+" => ["", ""],
      "-" => ["-", ""],
      "!" => ["\\lnot\\left(", "\\right)"]
    }.fetch(operator, [operator, ""])

    "#{left}#{latexify(statement)}#{right}"
  in SyntaxTree::VarRef[value: SyntaxTree::Const | SyntaxTree::Ident] |
     SyntaxTree::VCall[value: SyntaxTree::Const | SyntaxTree::Ident]
    latexify(node.value)
  else
    raise ArgumentError, "Using syntax not understood by latexify: #{PP.pp(node, +"")}"
  end
end

################################################################################
# Below here are some examples that mirror that ones in the python version.    #
################################################################################

def solve(a, b, c)
  (-b + Math.sqrt(b**2 - 4*a*c)) / (2*a)
end

puts latexify(find_method(method(:solve)))
# => \mathrm{solve}(a, b, c)\triangleq \frac{-b + \sqrt{b^{2} - 4ac}}{2a}

def sinc(x)
  if x == 0
    1
  else
    Math.sin(x) / x
  end
end

puts latexify(find_method(method(:sinc)))
# => \mathrm{sinc}(x)\triangleq \left\{ \begin{array}{ll} 1, & \mathrm{if} \ x = 0 \\ \frac{\sin{\left({x}\right)}}{x}, & \mathrm{otherwise} \end{array} \right.

def fib(x)
  if x == 0
    1
  elsif x == 1
    1
  else
    fib(x - 1) + fib(x - 2)
  end
end

puts latexify(find_method(method(:fib)))
# => \mathrm{fib}(x)\triangleq \left\{ \begin{array}{ll} 1, & \mathrm{if} \ x = 0 \\ 1, & \mathrm{if} \ x = 1 \\ \mathrm{fib}\left(x - 1\right) + \mathrm{fib}\left(x - 2\right), & \mathrm{otherwise} \end{array} \right.

def greek(alpha, beta, gamma, omega)
  alpha * beta + Math.gamma(gamma) + omega
end

puts latexify(find_method(method(:greek)))
# => \mathrm{greek}({\alpha}, {\beta}, {\gamma}, {\omega})\triangleq {\alpha}{\beta} + \Gamma\left({{\gamma}}\right) + {\omega}