/******************************************************************************

 * Top contributors (to current version):

 *   Abdalrhman Mohamed, Andrew Reynolds

 *

 * This file is part of the cvc5 project.

 *

 * Copyright (c) 2009-2021 by the authors listed in the file AUTHORS

 * in the top-level source directory and their institutional affiliations.

 * All rights reserved.  See the file COPYING in the top-level source

 * directory for licensing information.

 * ****************************************************************************

 *

 * Utility for reconstructing terms to match a grammar.

 */

#include "cvc5_private.h"

#ifndef CVC5__THEORY__QUANTIFIERS__SYGUS_RECONSTRUCT_H

#define CVC5__THEORY__QUANTIFIERS__SYGUS_RECONSTRUCT_H

#include <map>

#include <vector>

#include "expr/match_trie.h"

#include "smt/env_obj.h"

#include "theory/quantifiers/sygus/rcons_obligation.h"

#include "theory/quantifiers/sygus/rcons_type_info.h"

namespace cvc5 {

namespace theory {

namespace quantifiers {

using BuiltinSet = std::unordered_set<Node>;

using TypeBuiltinSetMap = std::unordered_map<TypeNode, BuiltinSet>;

using NodePairMap = std::unordered_map<Node, Node>;

/** SygusReconstruct

 *

 * This class implements an algorithm for reconstructing a given term such that

 * the reconstructed term is equivalent to the original term and its syntax

 * matches a specified grammar.

 *

 * Goal: find a term g in sygus type T0 that is equivalent to builtin term t0.

 *

 * rcons(t0, T0) returns g

 * {

 *   Obs: A set of pairs to reconstruct into T, where each pair is of the form

 *        (k, ts), where k is a skolem of a sygus type T and ts is a set of

 *        builtin terms of the type encoded by T.

 *

 *   Terms: A map from skolems k to a set of builtin terms in the pairs (k, ts).

 *          That is, Terms[k] = ts.

 *

 *   Sol: A map from skolems k to (possibly null) sygus terms of type T

 *        representing solutions to the obligations.

 *

 *   CandSols : A map from a skolem k to a set of possible solutions for its

 *              corresponding obligation. Whenever there is a successful match,

 *              it is added to this set.

 *

 *   Pool : A map from a sygus type T to a set of enumerated terms in T.

 *          The terms in this pool are patterns whose builtin analogs are used

 *          for matching against the terms to reconstruct ts in (k, ts).

 *

 *   let k0 be a fresh skolem of sygus type T0

 *   Obs[T0] += (k0, {t0})

 *

 *   TermsToRecons[T0] = {t0}

 *

 *   while Sol[k0] == null

 *     // map from T to terms pending addition to TermsToRecons

 *     TermsToRecons' = {}

 *     // enumeration phase

 *     for each subfield type T of T0

 *       // enumerated terms may contain variables zs ranging over all terms of

 *       // their type (subfield types of T0)

 *       s[zs] = nextEnum(T)

 *       if (s[zs] is ground)

 *         builtin = rewrite(toBuiltIn(s[zs]))

 *         // let X be the theory the solver is invoked with

 *         find (k, ts) in Obs s.t. |=_X ts[0] = builtin

 *         if no such pair exists

 *           k = newVar(T)

 *           ts = {builtin}

 *           Obs += (k, ts)

 *         markSolved((k, ts), s[zs])

 *       else if no s' in Pool[T] and matcher sigma s.t.

 *             rewrite(toBuiltIn(s')) * sigma = rewrite(toBuiltIn(s[zs]))

 *         Pool[T] += s[zs]

 *         for each t in TermsToRecons[T]

 *           TermsToRecons' += matchNewObs(t, s[zs])

 *     // match phase

 *     while TermsToRecons' != {}

 *       TermsToRecons'' = {}

 *       for each subfield type T of T0

 *         for each t in TermsToRecons'[T]

 *           TermsToRecons'[T] += t

 *           for each s[zs] in Pool[T]

 *             TermsToRecons'' += matchNewObs(t, s[zs])

 *         TermsToRecons' = TermsToRecons''

 *   g = Sol[k0]

 *   instantiate free variables of g with arbitrary sygus datatype values

 * }

 *

 * matchNewObs(t, s[zs]) returns TermsToRecons'

 * {

 *   u = rewrite(toBuiltIn(s[zs]))

 *   if match(u, t) == {toBuiltin(zs) -> sts}

 *     Sub = {} // substitution map from zs to corresponding new vars ks

 *     for each (z, st) in {zs -> sts}

 *       // let X be the theory the solver is invoked with

 *       if exists (k, ts) in Obs s.t. !=_X ts[0] = st

 *         ts += st

 *         Sub[z] = k

 *       else

 *         sk = newVar(typeOf(z))

 *         Sub[z] = sk

 *         TermsToRecons'[typeOf(z)] += st

 *     if Sol[sk] != null forall (z, sk) in Sub

 *       markSolved(k, s{Sub})

 *     else

 *       CandSol[k] += s{Sub}

 * }

 *

 * markSolved((k, ts), s)

 * {

 *   if Sol[k] != null

 *     return

 *   Sol[k] = s

 *   CandSols[k] += s

 *   Stack = {k}

 *   while Stack != {}

 *     k' = pop(Stack)

 *     for all k'' in CandSols keys

 *       for all s'[k'] in CandSols[k'']

 *         s'{k' -> Sol[k']}

 *         if s' does not contain free variables in Obs

 *           Sol[k''] = s'

 *           push(Stack, k'')

 * }

 */

{

 public:

  /**

   * Constructor.

   *

   * @param env reference to the environment

   * @param tds database for sygus terms

   * @param s statistics managed for the synth engine

   */

  SygusReconstruct(Env& env, TermDbSygus* tds, SygusStatistics& s);

  /** Reconstruct solution.

   *

   * Return (if possible) a sygus encoding of a node that is equivalent to sol

   * and whose syntax matches the grammar corresponding to sygus datatype stn.

   *

   * For example, given:

   *   sol = (* 2 x)

   *   stn = A sygus datatype encoding the grammar

   *           Start -> (c_PLUS Start Start) | c_x | c_0 | c_1

   * This method may return (c_PLUS c_x c_x) and set reconstructed to 1.

   *

   * @param sol the target term

   * @param stn the sygus datatype type encoding the syntax restrictions

   * @param reconstructed the flag to update, set to 1 if we successfully return

   *                      a node, otherwise it is set to -1

   * @param enumLimit a value to limit the effort spent by this class (roughly

   *                  equal to the number of intermediate terms to try)

   * @return the reconstructed sygus term

   */

  Node reconstructSolution(Node sol,

                           TypeNode stn,

                           int8_t& reconstructed,

                           uint64_t enumLimit);

 private:

  /** Match builtin term `t` with pattern `sz`.

   *

   * This function matches the builtin term to reconstruct `t` with the builtin

   * analog of the pattern `sz`. If the match succeeds, `sz` is added to the set

   * of candidate solutions for the obligation `ob` corresponding to the builtin

   * term `t` and a set of new sub-terms to reconstruct is returned. If there

   * are no new sub-terms to reconstruct, then `sz` is considered a solution to

   * obligation `ob` and `markSolved(ob, sz)` is called. For example, given:

   *

   * Obs = [(k0, {(+ 1 1)})]

   * Pool[T0] = {(c_+ z1 z2)}

   * CandSols = {}

   *

   * Then calling `matchNewObs((+ 1 1), (c_+ z1 z2))` will result in:

   *

   * Obs = [(k0, {(+ 1 1)}), (k1, {1})]

   * Pool[T0] = {(c_+ z1 z2)}

   * CandSols = {k0 -> {(c_+ k1 k1)}}

   *

   * and will return `{T0 -> {1}}`.

   *

   * Notice that the patterns/shapes in pool are generic, and thus, contain vars

   * `z` as opposed to skolems `k`. Also, notice that `(c_+ z1 z2)` is

   * instantiated with only one skolem `k1`, which is then added to the set of

   * obligations. That's because the builtin analogs of `z1` and `z2` match with

   * the same builtin term `1`.

   *

   * @param t builtin term we need to reconstruct

   * @param sz a pattern to match against `t`

   * @return a set of new builtin terms to reconstruct if the match succeeds

   */

  TypeBuiltinSetMap matchNewObs(Node t, Node sz);

  /** mark obligation `ob` as solved.

   *

   * This function first marks `s` as the complete/constant solution for

   * `ob`. Then it substitutes all instances of `ob` in partial solutions to

   * other obligations with `s`. The function repeats those two steps for each

   * partial solution that gets completed because of the second step. For

   * example, given:

   *

   * CandSols = {

   *   k0 -> {(c_+ k1 c_1)},

   *   k1 -> {(c_* k2 c_x)},

   *   k2 -> {c_2}

   * }

   * Sol = {}

   *

   * Then calling `markSolved(k2, c_2)` will result in:

   *

   * CandSols = {

   *   k0 -> {(c_+ k1 c_1), (c_+ (c_* c_2 c_x) c_1)},

   *   k1 -> {(c_* k2 c_x), (c_* c_2 c_x)},

   *   k2 -> {c_2}

   * }

   * Sol = {

   *   k0 -> (c_+ (c_* c_2 c_x) c_1),

   *   k1 -> (c_* c_2 c_x),

   *   k2 -> c_2

   * }

   *

   * @param ob free var to mark as solved and substitute

   * @param sol solution to `ob`

   */

  void markSolved(RConsObligation* ob, Node s);

  /**

   * Initialize a sygus enumerator and a candidate rewrite database for each

   * sygus datatype type.

   *

   * @param stn The sygus datatype type encoding the syntax restrictions

   */

  void initialize(TypeNode stn);

  /**

   * Remove reconstructed terms from the given set of terms to reconstruct.

   *

   * @param termsToRecons a set of terms to reconstruct possibly containing

   *                      already reconstructed ones

   */

  void removeReconstructedTerms(TypeBuiltinSetMap& termsToRecons);

  /**

   * Replace all variables in `n` with ground values. Before, calling `match`,

   * `matchNewObs` rewrites the builtin analog of `n` which contains variables.

   * The rewritten term, however, may contain fewer variables:

   *

   * rewrite((ite true z1 z2)) = z1 // n = (c_ite c_true z1 z2)

   *

   * In such cases, `matchNewObs` replaces the remaining variables (`z1`) with

   * obligations and ignores the eliminated ones (`z2`). Since the values of the

   * eliminated variables do not matter, they are replaced with some ground

   * values by calling this method.

   *

   * @param n a term containing variables

   * @return `n` with all vars in `n` replaced with ground values

   */

  Node mkGround(Node n) const;

  /**

   * A notification that s is equal to n * { vars -> subs }. This function

   * should return false if we do not wish to be notified of further matches.

   */

  bool notify(Node s,

              Node n,

              std::vector<Node>& vars,

              std::vector<Node>& subs) override;

  /**

   * Reset the state of this SygusReconstruct object.

   */

  void clear();

  /**

   * Print the pool of patterns/shape used in the matching phase.

   *

   * \note requires enabling "sygus-rcons" trace

   *

   * @param pool a pool of patterns/shapes to print

   */

  void printPool(

      const std::unordered_map<TypeNode, std::vector<Node>>& pool) const;

  /** pointer to the sygus term database */

  TermDbSygus* d_tds;

  /** reference to the statistics of parent */

  SygusStatistics& d_stats;

  /** a list of obligations to solve */

  std::vector<std::unique_ptr<RConsObligation>> d_obs;

  /** a map from a sygus datatype type to its reconstruction info */

  std::unordered_map<TypeNode, RConsTypeInfo> d_stnInfo;

  /** a map from an obligation's skolem to its sygus solution (if it exists) */

  std::unordered_map<TNode, TNode> d_sol;

  /** a map from a candidate solution to its sub-obligations */

  std::unordered_map<Node, std::vector<RConsObligation*>> d_subObs;

  /** a map from a candidate solution to its parent obligation */

  std::unordered_map<Node, RConsObligation*> d_parentOb;

  /** a cache of sygus variables treated as ground terms by matching */

  std::unordered_map<Node, Node> d_sygusVars;

  /** A trie for filtering out redundant terms from the paterns pool */

  expr::MatchTrie d_poolTrie;

};

}  // namespace quantifiers

}  // namespace theory

}  // namespace cvc5

#endif  // CVC5__THEORY__QUANTIFIERS__SYGUS_RECONSTRUCT_H