shot_ds

defmodule ShotDs.Util.Lexer do
  @moduledoc """
  Contains a `tokenize/1` function for tokenizing a string representing a
  formula in THF syntax or a TPTP TH0/TH1 problem file using `NimbleParsec`.
  This is mainly used as a preprocessing step for parsing. For information about
  the returned structure, see
  https://hexdocs.pm/nimble_parsec/NimbleParsec.html.

  Every token is a 3-tuple `{type, value, byte_offset}` where `byte_offset` is
  the position of the first byte of the token in the original input string.
  Downstream parsers use this to report the position of type errors.

  ## Examples

      iex> {:ok, tokens, "", _, _, _} = tokenize("A & B")
      iex> tokens
      [{:var "A", 0}, {:and, "&", 2}, {:var, "B", 4}]
  """

  import NimbleParsec

  @typedoc """
  Tokens are generated by the lexer as a list of `{type, value, byte_offset}`
  triples, where `type` is a label given as atom, `value` is the original
  string representation of the token, and `byte_offset` is the position in
  the original input.
  """
  @type tokens() :: [{atom(), String.t() | atom(), non_neg_integer()}]

  whitespace = ascii_string([?\s, ?\t, ?\n, ?\r], min: 1) |> ignore()

  comment = string("%") |> concat(ascii_string([not: ?\n], min: 0)) |> ignore()

  @keywords %{
    "include" => :include,
    "thf" => :thf
  }

  @roles %{
    "type" => :type,
    "axiom" => :axiom,
    "definition" => :definition,
    "conjecture" => :conjecture,
    "negated_conjecture" => :negated_conjecture,
    "lemma" => :lemma,
    "hypothesis" => :hypothesis,
    "assumption" => :assumption
  }

  @doc false
  def categorize_atom(name) do
    cond do
      Map.has_key?(@keywords, name) -> {:keyword, @keywords[name]}
      Map.has_key?(@roles, name) -> {:role, @roles[name]}
      true -> {:atom, name}
    end
  end

  @doc false
  def attach_offset(rest, [{type, value}], context, _line, end_offset) do
    start = end_offset - byte_size(to_string(value))
    {rest, [{type, value, start}], context}
  end

  system_symbol =
    string("$")
    |> ascii_string([?a..?z, ?A..?Z, ?0..?9, ?_], min: 1)
    |> reduce({Enum, :join, []})
    |> unwrap_and_tag(:system)
    |> post_traverse({__MODULE__, :attach_offset, []})

  variable =
    ascii_string([?A..?Z], 1)
    |> optional(ascii_string([?a..?z, ?A..?Z, ?0..?9, ?_], min: 1))
    |> reduce({Enum, :join, []})
    |> unwrap_and_tag(:var)
    |> post_traverse({__MODULE__, :attach_offset, []})

  atom_ident =
    ascii_string([?a..?z], 1)
    |> optional(ascii_string([?a..?z, ?A..?Z, ?0..?9, ?_], min: 1))
    |> reduce({Enum, :join, []})
    |> map({__MODULE__, :categorize_atom, []})
    |> post_traverse({__MODULE__, :attach_offset, []})

  distinct_object =
    ignore(string("'"))
    |> concat(ascii_string([not: ?'], min: 1))
    |> ignore(string("'"))
    |> unwrap_and_tag(:distinct)
    |> post_traverse({__MODULE__, :attach_offset, []})

  symbols =
    choice([
      # Connectives
      string("<=>") |> replace({:equiv, "<=>"}),
      string("=>") |> replace({:implies, "=>"}),
      string("<=") |> replace({:implied_by, "<="}),
      string("<~>") |> replace({:xor, "<~>"}),
      string("~|") |> replace({:nor, "~|"}),
      string("~&") |> replace({:nand, "~&"}),

      # Quantifiers and Operators
      string("!!") |> replace({:forall, "!!"}),
      string("!=") |> replace({:neq, "!="}),
      string("!>") |> replace({:forall_type, "!>"}),
      string("!") |> replace({:forall, "!"}),
      string("??") |> replace({:exists, "??"}),
      string("?") |> replace({:exists, "?"}),
      string("^") |> replace({:lambda, "^"}),
      string("@") |> replace({:app, "@"}),
      string("~") |> replace({:not, "~"}),
      string("|") |> replace({:or, "|"}),
      string("&") |> replace({:and, "&"}),
      string("=") |> replace({:eq, "="}),

      # Punctuation
      string("(") |> replace({:lparen, "("}),
      string(")") |> replace({:rparen, ")"}),
      string("[") |> replace({:lbracket, "["}),
      string("]") |> replace({:rbracket, "]"}),
      string(":") |> replace({:colon, ":"}),
      string(",") |> replace({:comma, ","}),
      string(".") |> replace({:dot, "."}),

      # Types
      string(">") |> replace({:arrow, ">"})
    ])
    |> post_traverse({__MODULE__, :attach_offset, []})

  defparsec(
    :tokenize,
    repeat(
      choice([
        whitespace,
        comment,
        symbols,
        system_symbol,
        variable,
        atom_ident,
        distinct_object
      ])
    )
  )

  @doc """
  Converts a byte offset into the source string into a `{line, column}` pair
  (1-indexed). Used for decorating type-error messages produced by the parser.
  """
  @spec byte_offset_to_line_col(String.t(), non_neg_integer()) ::
          {pos_integer(), pos_integer()}
  def byte_offset_to_line_col(source, target_offset) do
    prefix = binary_part(source, 0, min(target_offset, byte_size(source)))
    lines = String.split(prefix, "\n")
    line = length(lines)
    last_line = List.last(lines, "")
    {line, String.length(last_line) + 1}
  end
end