zee3_marco

defmodule Zee3.MARCO do
  @moduledoc """
  Documentation for `Zee3.MARCO` algorithm, which implements
  the [MARCO argoritm](https://markliffiton.com/marco/)
  for Zee3-based constraints.
  """

  alias Zee3.Smt2
  alias Zee3.Solver
  alias Zee3.StdLib.Sort

  alias Zee3.MARCO.State

  @doc """
  Enumerate the all the maximum satisfying sets (MSSs)
  and minimum unsatisfiable cores (MUSs) for a set of
  constraints. Returns a Stream that yields tuples of
  the form `{:mss, set}` or `{:mus, set}`.

  > #### Warning {: .warning}
  >
  > The order of the returned sets is non-deterministic
  """
  @spec enumerate_sets(map(), (pid() -> :ok)) :: Enumerable.t({:mss | :mus, MapSet.t()})
  def enumerate_sets(constraint_map, background \\ fn _pid -> :ok end)
      when is_map(constraint_map) do

    start_fun = fn ->
      # Start a new Z3 solver process
      {:ok, subset_solver_pid} = Solver.start_link()
      # Add the background assumptions to the solver
      :ok = background.(subset_solver_pid)
      setup(subset_solver_pid, constraint_map)
    end

    next_fun = fn state ->
      case next_set(state) do
        {:halt, new_state} ->
          {:halt, new_state}

        {set, new_state} ->
          {[set], new_state}
      end
    end

    after_fun = fn state ->
      Solver.stop(state.subset_solver)
      Solver.stop(state.map_solver)
      :ok
    end

    Stream.resource(start_fun, next_fun, after_fun)
  end

  @doc """
  Enumerate the MUSs for this set of constraints. Returns a Stream.

  > #### Warning {: .warning}
  >
  > The order of the returned MUSs is non-deterministic
  """
  @spec enumerate_mus(map(), (pid() -> :ok)) :: Enumerable.t(MapSet.t())
  def enumerate_mus(constraint_map, background \\ fn pid -> pid end) do
    constraint_map
    |> enumerate_sets(background)
    |> Stream.filter(fn {:mus, _set} -> true; _ -> false end)
    |> Stream.map(fn {:mus, set} -> set end)
  end

  @doc """
  Enumerate the MSSs for this set of constraints. Returns a Stream.

  > #### Warning {: .warning}
  >
  > The order of the returned MSSs is non-deterministic
  """
  @spec enumerate_mss(map(), (pid() -> :ok)) :: Enumerable.t(MapSet.t())
  def enumerate_mss(constraint_map, background \\ fn pid -> pid end) do
    constraint_map
    |> enumerate_sets(background)
    |> Stream.filter(fn {:mss, _set} -> true; _ -> false end)
    |> Stream.map(fn {:mss, set} -> set end)
  end

  @doc """
  Returns all maximum satisfying sets (MSSs)
  and minimum unsatisfiable cores (MUSs) for a set of
  constraints. Returns a list of tuples of
  the form `{:mss, set}` or `{:mus, set}`.

  This function eagerly evaluates all sets, which can take
  a rather long time.
  If you need control over how long to wait until you quit
  generating more sets or if you want to stream sets as they
  are generated, use the `enumerate_sets/2` function.

  The results are sorted by set (generally, smaller sets come first)
  so that the result of this function is deterministic.
  """
  def all_sets(constraint_map, background \\ fn pid -> pid end) do
    constraint_map
    |> enumerate_sets(background)
    |> Enum.into([])
    |> Enum.sort_by(fn {_, set} -> set end)
  end

  @doc """
  Returns all minimum unsatisfiable cores (MUSs)
  for a set of constraints. Returns a list.

  This function eagerly evaluates all MUSs, which can take
  a rather long time.
  If you need control over how long to wait until you quit
  generating more sets or if you want to stream sets as they
  are generated, use the `enumerate_mus/2` function.

  The results are sorted by set (generally, smaller sets come first)
  so that the result of this function is deterministic.
  """
  def all_mus(constraint_map, background \\ fn pid -> pid end) do
    constraint_map
    |> enumerate_mus(background)
    |> Enum.into([])
    |> Enum.sort()
  end

  @doc """
  Returns all minimum unsatisfiable cores (MSSs)
  for a set of constraints. Returns a list.

  This function eagerly evaluates all MSSs, which can take
  a rather long time.
  If you need control over how long to wait until you quit
  generating more sets or if you want to stream sets as they
  are generated, use the `enumerate_mus/2` function.

  The results are sorted by set (generally, smaller sets come first)
  so that the result of this function is deterministic.
  """
  def all_mss(constraint_map, background \\ fn pid -> pid end) do
    constraint_map
    |> enumerate_mss(background)
    |> Enum.into([])
    |> Enum.sort()
  end

  @spec setup(pid(), %{binary() => Smt2.smt_like()}) :: State.t()
  defp setup(subset_solver_pid, constraints_map) when is_map(constraints_map) do
    {:ok, map_solver_pid} = Solver.start_link()

    constraints =
      constraints_map
      |> Map.keys()
      |> MapSet.new()

    for {c_id, _c} <- constraints_map do
      Solver.declare_const(map_solver_pid, c_id, Sort.bool())
      Solver.declare_const(subset_solver_pid, c_id, Sort.bool())
    end

    for {c_id, c} <- constraints_map do
      Solver.assert(
        subset_solver_pid,
        Smt2.call("implies", [Smt2.symbol(c_id), c])
      )
    end

    %State{
      constraints_map: constraints_map,
      constraints: constraints,
      seed: MapSet.new(),
      subset_solver: subset_solver_pid,
      map_solver: map_solver_pid
    }
  end

  @spec grow(State.t()) :: State.t()
  defp grow(state) do
    constraints_minus_seed =
      MapSet.difference(
        state.constraints,
        state.seed
      )

    mss =
      Enum.reduce(constraints_minus_seed, state.seed, fn c, seed ->
        new_seed = MapSet.put(state.seed, c)
        if is_satisfiable?(state.subset_solver, new_seed) do
          new_seed
        else
          seed
        end
      end)

    {mss, %{state | seed: mss}}
  end

  @spec shrink(State.t()) :: State.t()
  defp shrink(state) do
    mus =
      Enum.reduce(state.seed, state.seed, fn c, seed ->
        new_seed = MapSet.delete(state.seed, c)
        if is_satisfiable?(state.subset_solver, new_seed) do
          seed
        else
          new_seed
        end
      end)

    {mus, %{state | seed: mus}}
  end

  @spec block_up(State.t()) :: State.t()
  defp block_up(state) do
    new_map_constraints =
      state.seed
      |> Enum.map(fn c_id -> Smt2.call("not", [Smt2.symbol(c_id)]) end)
      |> then(fn args -> Smt2.call("or", args) end)

    Solver.assert(
      state.map_solver,
      new_map_constraints
    )

    state
  end

  @spec block_down(State.t()) :: State.t()
  defp block_down(state) do
    new_map_constraints =
      state.constraints
      |> MapSet.difference(state.seed)
      |> Enum.map(&Smt2.symbol/1)
      |> then(fn args -> Smt2.call("or", args) end)

    Solver.assert(
      state.map_solver,
      new_map_constraints
    )

    state
  end

  @spec next_set(State.t()) :: {{:mss | :mus, State.t()}, State.t()} | :halt
  defp next_set(state) do
    case fetch_unexplored_seed(state) do
      {:ok, seed} ->
        # Update the seed in the state
        state = %{state | seed: seed}

        # The `fetch_unexplored_seed/1` function has found a
        # seed that works on the map solver, which means it can
        # be a potentially valid solution, BUT we need to check
        # that into the solver which actually knows the content
        # of the constraints.
        #
        # For this, we switch to the `subset_solver`.
        if is_satisfiable?(state.subset_solver, seed) do
          # The constraint is satisfiable, which means we
          # have found a possible maximally satisfying set (MSS).
          # We will run the following sequence of steps:

          # 1. Grow the seed as far as we can by walking up the
          #    powerset lattice in a greedy way, always moving
          #    away from the empty set
          {mss, state} = grow(state)

          # 2. Now that we know that this is satisfiable, we
          #    know that all subsets that sit below this one
          #    in the powerset lattice are also satisfiable,
          #    so we have to block the solver from exploring
          #    them (that's why the function is called `block_down/1`)
          state = block_down(state)

          # 3. Return the MSS we have found and update the seed
          #    in the state; the new seed is the MSS we've just found.
          {{:mss, mss}, %{state | seed: mss}}
        else

          {mus, state} = shrink(state)
          state = block_up(state)
          {{:mus, mus}, %{state | seed: mus}}
        end

      :error ->
        # There are no more valid seed that we haven't explored yet.
        # The map solver is telling us that the powerset lattice
        # has been fully explored and there is no need to look
        # for any more seeds.
        #
        # We have used the magic of SAT solvers, the mathematics
        # of set theory and the properties of walking up and down
        # a lattice to effectively explore a very large search space
        # of 2^n size, with only a small fraction of the work required
        # to explore the full search space.
        #
        # We can now return the final state.
        {:halt, state}
    end
  end

  @spec fetch_unexplored_seed(State.t()) :: {:ok, MapSet.t()} | :error
  defp fetch_unexplored_seed(state) do
    case Zee3.check_sat_and_get_model!(state.map_solver) do
      {:sat, model} ->
        new_seed =
          Enum.reduce(state.constraints, MapSet.new(), fn c, new_seed ->
            case Map.fetch(model, c) do
              {:ok, false} ->
                new_seed

              _other ->
                MapSet.put(new_seed, c)
            end
          end)

        {:ok, new_seed}

      :unsat ->
        :error

      :unknown ->
        :error
    end
  end

  defp seed_to_symbols(seed) do
    Enum.map(seed, &Smt2.symbol/1)
  end

  defp is_satisfiable?(subset_solver_pid, seed) do
    symbols = seed_to_symbols(seed)

    case Solver.check_sat_assuming!(subset_solver_pid, symbols) do
      :sat -> true
      :unsat -> false
      :unknown -> false
    end
  end
end