shot_ds

defmodule ShotDs.Util.Lexer do
  @moduledoc """
  Contains a `tokenize/1` function for tokenizing a string representing a
  formula in TH0 syntax or a TPTP TH0 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.

  ## Examples

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

  import NimbleParsec

  @typedoc """
  Tokens are generated by the lexer as a list which contains `{type, data}`
  pairs, where `type` is a label given as atom and `data` is the original
  string representation of the token.
  """
  @type tokens() :: [{atom(), String.t()}]

  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,
    "lemma" => :lemma,
    "hypothesis" => :hypothesis
  }

  @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

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

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

  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, []})

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

  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({:pi, "!!"}),
      string("!=") |> replace({:neq, "!="}),
      string("!") |> replace({:forall, "!"}),
      string("??") |> replace({:sigma, "??"}),
      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, ">"})
    ])

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