text: /*
   * Souffle - A Datalog Compiler
   * Copyright (c) 2018, The Souffle Developers. All rights reserved.
   * Licensed under the Universal Permissive License v 1.0 as shown at:
   * - https://opensource.org/licenses/UPL
   * - <souffle root>/licenses/SOUFFLE-UPL.txt
   */
  
  /************************************************************************
   *
   * @file InlineRelations.cpp
   *
   * Define classes and functionality related to inlining.
   *
   ***********************************************************************/
  
  #include "ast/transform/InlineRelations.h"
  #include "AggregateOp.h"
  #include "RelationTag.h"
  #include "ast/Aggregator.h"
  #include "ast/Argument.h"
  #include "ast/Atom.h"
  #include "ast/BinaryConstraint.h"
  #include "ast/BooleanConstraint.h"
  #include "ast/Clause.h"
  #include "ast/Constant.h"
  #include "ast/Constraint.h"
  #include "ast/Functor.h"
  #include "ast/IntrinsicFunctor.h"
  #include "ast/Literal.h"
  #include "ast/Negation.h"
  #include "ast/Node.h"
  #include "ast/NumericConstant.h"
  #include "ast/Program.h"
  #include "ast/QualifiedName.h"
  #include "ast/RecordInit.h"
  #include "ast/Relation.h"
  #include "ast/TranslationUnit.h"
  #include "ast/TypeCast.h"
  #include "ast/UnnamedVariable.h"
  #include "ast/UserDefinedFunctor.h"
  #include "ast/Variable.h"
  #include "ast/analysis/PolymorphicObjects.h"
  #include "ast/utility/NodeMapper.h"
  #include "ast/utility/Utils.h"
  #include "ast/utility/Visitor.h"
  #include "souffle/BinaryConstraintOps.h"
  #include "souffle/utility/MiscUtil.h"
  #include <algorithm>
  #include <cassert>
  #include <cstddef>
  #include <memory>
  #include <optional>
  #include <ostream>
  #include <set>
  #include <string>
  #include <utility>
  #include <vector>
  
  namespace souffle::ast::transform {
  
  template <class T>
  class NullableVector {
  private:
      std::vector<T> vector;
      bool valid = false;
  
  public:
      NullableVector() = default;
      NullableVector(std::vector<T> vector) : vector(std::move(vector)), valid(true) {}
  
      bool isValid() const {
          return valid;
      }
  
      const std::vector<T>& getVector() const {
          assert(valid && "Accessing invalid vector!");
          return vector;
      }
  };
  
  NullableVector<std::vector<Literal*>> getInlinedLiteral(Program&, Literal*);
  
  /**
   * Replace constants in the head of inlined clauses with (constrained) variables.
   */
  bool normaliseInlinedHeads(Program& program) {
      bool changed = false;
      static int newVarCount = 0;
  
      // Go through the clauses of all inlined relations
      for (Relation* rel : program.getRelations()) {
          if (!rel->hasQualifier(RelationQualifier::INLINE)) {
              continue;
          }
  
          for (Clause* clause : getClauses(program, *rel)) {
              // Set up the new clause with an empty body and no arguments in the head
              auto newClause = mk<Clause>(clause->getHead()->getQualifiedName(), clause->getSrcLoc());
              newClause->setBodyLiterals(souffle::clone(clause->getBodyLiterals()));
              auto clauseHead = newClause->getHead();
  
              // Set up the head arguments in the new clause
              for (Argument* arg : clause->getHead()->getArguments()) {
                  if (auto* constant = as<Constant>(arg)) {
                      // Found a constant in the head, so replace it with a variable
                      std::stringstream newVar;
                      newVar << "<new_var_" << newVarCount++ << ">";
                      clauseHead->addArgument(mk<ast::Variable>(newVar.str()));
  
                      // Add a body constraint to set the variable's value to be the original constant
                      newClause->addToBody(mk<BinaryConstraint>(BinaryConstraintOp::EQ,
                              mk<ast::Variable>(newVar.str()), souffle::clone(constant)));
                  } else {
                      // Already a variable
                      clauseHead->addArgument(souffle::clone(arg));
                  }
              }
  
              // Replace the old clause with this one
              program.addClause(std::move(newClause));
              program.removeClause(clause);
              changed = true;
          }
      }
  
      return changed;
  }
  
  /**
   * Removes all underscores in all atoms of inlined relations
   */
  bool nameInlinedUnderscores(Program& program) {
      struct M : public NodeMapper {
          mutable bool changed = false;
          const std::set<QualifiedName> inlinedRelations;
          bool replaceUnderscores;
  
          M(std::set<QualifiedName> inlinedRelations, bool replaceUnderscores)
                  : inlinedRelations(std::move(inlinedRelations)), replaceUnderscores(replaceUnderscores) {}
  
          Own<Node> operator()(Own<Node> node) const override {
              static int underscoreCount = 0;
  
              if (!replaceUnderscores) {
                  // Check if we should start replacing underscores for this node's subnodes
                  if (auto* atom = as<Atom>(node)) {
                      if (inlinedRelations.find(atom->getQualifiedName()) != inlinedRelations.end()) {
                          // Atom associated with an inlined relation, so replace the underscores
                          // in all of its subnodes with named variables.
                          M replace(inlinedRelations, true);
                          node->apply(replace);
                          changed |= replace.changed;
                          return node;
                      }
                  }
              } else if (isA<UnnamedVariable>(node)) {
                  // Give a unique name to the underscored variable
                  // TODO (azreika): need a more consistent way of handling internally generated variables in
                  // general
                  std::stringstream newVarName;
                  newVarName << "<underscore_" << underscoreCount++ << ">";
                  changed = true;
                  return mk<ast::Variable>(newVarName.str());
              }
  
              node->apply(*this);
              return node;
          }
      };
  
      // Store the names of all relations to be inlined
      std::set<QualifiedName> inlinedRelations;
      for (Relation* rel : program.getRelations()) {
          if (rel->hasQualifier(RelationQualifier::INLINE)) {
              inlinedRelations.insert(rel->getQualifiedName());
          }
      }
  
      // Apply the renaming procedure to the entire program
      M update(inlinedRelations, false);
      program.apply(update);
      return update.changed;
  }
  
  /**
   * Checks if a given clause contains an atom that should be inlined.
   */
  bool containsInlinedAtom(const Program& program, const Clause& clause) {
      bool foundInlinedAtom = false;
  
      visit(clause, [&](const Atom& atom) {
          Relation* rel = getRelation(program, atom.getQualifiedName());
          if (rel->hasQualifier(RelationQualifier::INLINE)) {
              foundInlinedAtom = true;
          }
      });
  
      return foundInlinedAtom;
  }
  
  /**
   * Reduces a vector of substitutions.
   * Returns false only if matched argument pairs are found to be incompatible.
   */
  bool reduceSubstitution(std::vector<std::pair<Argument*, Argument*>>& sub) {
      // Keep trying to reduce the substitutions until we reach a fixed point.
      // Note that at this point no underscores ('_') or counters ('$') should appear.
      bool done = false;
      while (!done) {
          done = true;
  
          // Try reducing each pair by one step
          for (std::size_t i = 0; i < sub.size(); i++) {
              auto currPair = sub[i];
              Argument* lhs = currPair.first;
              Argument* rhs = currPair.second;
  
              // Start trying to reduce the substitution
              // Note: Can probably go further with this substitution reduction
              if (*lhs == *rhs) {
                  // Get rid of redundant `x = x`
                  sub.erase(sub.begin() + i);
                  done = false;
              } else if (isA<Constant>(lhs) && isA<Constant>(rhs)) {
                  // Both are constants but not equal (prev case => !=)
                  // Failed to unify!
                  return false;
              } else if (isA<RecordInit>(lhs) && isA<RecordInit>(rhs)) {
                  // Note: we will not deal with the case where only one side is
                  // a record and the other is a variable, as variables can be records
                  // on a deeper level.
                  std::vector<Argument*> lhsArgs = static_cast<RecordInit*>(lhs)->getArguments();
                  std::vector<Argument*> rhsArgs = static_cast<RecordInit*>(rhs)->getArguments();
  
                  if (lhsArgs.size() != rhsArgs.size()) {
                      // Records of unequal size can't be equated
                      return false;
                  }
  
                  // Equate all corresponding arguments
                  for (std::size_t i = 0; i < lhsArgs.size(); i++) {
                      sub.push_back(std::make_pair(lhsArgs[i], rhsArgs[i]));
                  }
  
                  // Get rid of the record equality
                  sub.erase(sub.begin() + i);
                  done = false;
              } else if ((isA<RecordInit>(lhs) && isA<Constant>(rhs)) ||
                         (isA<Constant>(lhs) && isA<RecordInit>(rhs))) {
                  // A record =/= a constant
                  return false;
              }
          }
      }
  
      return true;
  }
  
  /**
   * Returns the nullable vector of substitutions needed to unify the two given atoms.
   * If unification is not successful, the returned vector is marked as invalid.
   * Assumes that the atoms are both of the same relation.
   */
  NullableVector<std::pair<Argument*, Argument*>> unifyAtoms(Atom* first, Atom* second) {
      std::vector<std::pair<Argument*, Argument*>> substitution;
  
      std::vector<Argument*> firstArgs = first->getArguments();
      std::vector<Argument*> secondArgs = second->getArguments();
  
      // Create the initial unification equalities
      for (std::size_t i = 0; i < firstArgs.size(); i++) {
          substitution.push_back(std::make_pair(firstArgs[i], secondArgs[i]));
      }
  
      // Reduce the substitutions
      bool success = reduceSubstitution(substitution);
      if (success) {
          return NullableVector<std::pair<Argument*, Argument*>>(substitution);
      } else {
          // Failed to unify the two atoms
          return NullableVector<std::pair<Argument*, Argument*>>();
      }
  }
  
  /**
   * Inlines the given atom based on a given clause.
   * Returns the vector of replacement literals and the necessary constraints.
   * If unification is unsuccessful, the vector of literals is marked as invalid.
   */
  std::pair<NullableVector<Literal*>, std::vector<BinaryConstraint*>> inlineBodyLiterals(
          Atom* atom, Clause* atomInlineClause) {
      bool changed = false;
      std::vector<Literal*> addedLits;
      std::vector<BinaryConstraint*> constraints;
  
      // Rename the variables in the inlined clause to avoid conflicts when unifying multiple atoms
      // - particularly when an inlined relation appears twice in a clause.
      static int inlineCount = 0;
  
      // Make a temporary clone so we can rename variables without fear
      auto atomClause = souffle::clone(atomInlineClause);
  
      struct VariableRenamer : public NodeMapper {
          int varnum;
          VariableRenamer(int varnum) : varnum(varnum) {}
          Own<Node> operator()(Own<Node> node) const override {
              if (auto* var = as<ast::Variable>(node)) {
                  // Rename the variable
                  auto newVar = souffle::clone(var);
                  std::stringstream newName;
                  newName << "<inlined_" << var->getName() << "_" << varnum << ">";
                  newVar->setName(newName.str());
                  return newVar;
              }
              node->apply(*this);
              return node;
          }
      };
  
      VariableRenamer update(inlineCount);
      atomClause->apply(update);
  
      inlineCount++;
  
      // Get the constraints needed to unify th366Ge two atoms
      NullableVector<std::pair<Argument*, Argument*>> res = unifyAtoms(atomClause->getHead(), atom);
      if (res.isValid()) {
          changed = true;
          for (std::pair<Argument*, Argument*> pair : res.getVector()) {
              // FIXME: float equiv (`FEQ`)
              constraints.push_back(new BinaryConstraint(
                      BinaryConstraintOp::EQ, souffle::clone(pair.first), souffle::clone(pair.second)));
          }
  
          // Add in the body of the current clause of the inlined atom
          for (Literal* lit : atomClause->getBodyLiterals()) {
              // FIXME: This is a horrible hack.  Should convert
              // addedLits to hold Own<>
              addedLits.push_back(souffle::clone(lit).release());
          }
      }
  
      if (changed) {
          return std::make_pair(NullableVector<Literal*>(addedLits), constraints);
      } else {
          return std::make_pair(NullableVector<Literal*>(), constraints);
      }
  }
  
  /**
   * Returns the negated version of a given literal
   */
  Literal* negateLiteral(Literal* lit) {
      // FIXME: Ownership semantics
      if (auto* atom = as<Atom>(lit)) {
          auto* neg = new Negation(souffle::clone(atom));
          return neg;
      } else if (auto* neg = as<Negation>(lit)) {
          Atom* atom = souffle::clone(neg->getAtom()).release();
          return atom;
      } else if (auto* cons = as<Constraint>(lit)) {
          Constraint* newCons = souffle::clone(cons).release();
          negateConstraintInPlace(*newCons);
          return newCons;
      }
  
      fatal("unsupported literal type: %s", *lit);
  }
  
  /**
   * Return the negated version of a disjunction of conjunctions.
   * E.g. (a1(x) ^ a2(x) ^ ...) v (b1(x) ^ b2(x) ^ ...) --into-> (!a1(x) ^ !b1(x)) v (!a2(x) ^ !b2(x)) v ...
   */
  std::vector<std::vector<Literal*>> combineNegatedLiterals(std::vector<std::vector<Literal*>> litGroups) {
      std::vector<std::vector<Literal*>> negation;
  
      // Corner case: !() = ()
      if (litGroups.empty()) {
          return negation;
      }
  
      std::vector<Literal*> litGroup = litGroups[0];
      if (litGroups.size() == 1) {
          // !(a1 ^ a2 ^ a3 ^ ...) --into-> !a1 v !a2 v !a3 v ...
          for (auto& i : litGroup) {
              std::vector<Literal*> newVec;
              newVec.push_back(negateLiteral(i));
              negation.push_back(newVec);
          }
  
          // Done!
          return negation;
      }
  
      // Produce the negated versions of all disjunctions ignoring the first recursively
      std::vector<std::vector<Literal*>> combinedRHS = combineNegatedLiterals(
              std::vector<std::vector<Literal*>>(litGroups.begin() + 1, litGroups.end()));
  
      // We now just need to add the negation of a single literal from the untouched LHS
      // to every single conjunction on the RHS to finalise creating every possible combination
      for (Literal* lhsLit : litGroup) {
          for (std::vector<Literal*> rhsVec : combinedRHS) {
              std::vector<Literal*> newVec;
              newVec.push_back(negateLiteral(lhsLit));
  
              for (Literal* lit : rhsVec) {
                  // FIXME: This is a horrible hack.  Should convert
                  // newVec to hold Own<>
                  newVec.push_back(souffle::clone(lit).release());
              }
  
              negation.push_back(newVec);
          }
      }
  
      for (std::vector<Literal*> rhsVec : combinedRHS) {
          for (Literal* lit : rhsVec) {
              delete lit;
          }
      }
  
      return negation;
  }
  
  /**
   * Forms the bodies that will replace the negation of a given inlined atom.
   * E.g. a(x) <- (a11(x), a12(x)) ; (a21(x), a22(x)) => !a(x) <- (!a11(x), !a21(x)) ; (!a11(x), !a22(x)) ; ...
   * Essentially, produce every combination (m_1 ^ m_2 ^ ...) where m_i is the negation of a literal in the
   * ith rule of a.
   */
  std::vector<std::vector<Literal*>> formNegatedLiterals(Program& program, Atom* atom) {
      // Constraints added to unify atoms should not be negated and should be added to
      // all the final rule combinations produced, and so should be stored separately.
      std::vector<std::vector<Literal*>> addedBodyLiterals;
      std::vector<std::vector<BinaryConstraint*>> addedConstraints;
  
      // Go through every possible clause associated with the given atom
      for (Clause* inClause : getClauses(program, *getRelation(program, atom->getQualifiedName()))) {
          // Form the replacement clause by inlining based on the current clause
          std::pair<NullableVector<Literal*>, std::vector<BinaryConstraint*>> inlineResult =
                  inlineBodyLiterals(atom, inClause);
          NullableVector<Literal*> replacementBodyLiterals = inlineResult.first;
          std::vector<BinaryConstraint*> currConstraints = inlineResult.second;
  
          if (!replacementBodyLiterals.isValid()) {
              // Failed to unify, so just move on
              continue;
          }
  
          addedBodyLiterals.push_back(replacementBodyLiterals.getVector());
          addedConstraints.push_back(currConstraints);
      }
  
      // We now have a list of bodies needed to inline the given atom.
      // We want to inline the negated version, however, which is done using De Morgan's Law.
      std::vector<std::vector<Literal*>> negatedAddedBodyLiterals = combineNegatedLiterals(addedBodyLiterals);
  
      // Add in the necessary constraints to all the body literals
      for (auto& negatedAddedBodyLiteral : negatedAddedBodyLiterals) {
          for (std::vector<BinaryConstraint*> constraintGroup : addedConstraints) {
              for (BinaryConstraint* constraint : constraintGroup) {
                  // FIXME: This is a horrible hack.  Should convert
                  // negatedBodyLiterals to hold Own<>
                  negatedAddedBodyLiteral.push_back(souffle::clone(constraint).release());
              }
          }
      }
  
      // Free up the old body literals and constraints
      for (std::vector<Literal*> litGroup : addedBodyLiterals) {
          for (Literal* lit : litGroup) {
              delete lit;
          }
      }
      for (std::vector<BinaryConstraint*> consGroup : addedConstraints) {
          for (Constraint* cons : consGroup) {
              delete cons;
          }
      }
  
      return negatedAddedBodyLiterals;
  }
  
  /**
   * Renames all variables in a given argument uniquely.
   */
  void renameVariables(Argument* arg) {
      static int varCount = 0;
      varCount++;
  
      struct M : public NodeMapper {
          int varnum;
          M(int varnum) : varnum(varnum) {}
          Own<Node> operator()(Own<Node> node) const override {
              if (auto* var = as<ast::Variable>(node)) {
                  auto newVar = souffle::clone(var);
                  std::stringstream newName;
                  newName << var->getName() << "-v" << varnum;
                  newVar->setName(newName.str());
                  return newVar;
              }
              node->apply(*this);
              return node;
          }
      };
  
      M update(varCount);
      arg->apply(update);
  }
  
  // Performs a given binary op on a list of aggregators recursively.
  // E.g. ( <aggr1, aggr2, aggr3, ...>, o > = (aggr1 o (aggr2 o (agg3 o (...))))
  // TODO (azreika): remove aggregator support
  Argument* combineAggregators(std::vector<Aggregator*> aggrs, std::string fun) {
      // Due to variable scoping issues with aggregators, we rename all variables uniquely in the
      // added aggregator
      renameVariables(aggrs[0]);
  
      if (aggrs.size() == 1) {
          return aggrs[0];
      }
  
      Argument* rhs = combineAggregators(std::vector<Aggregator*>(aggrs.begin() + 1, aggrs.end()), fun);
  
      Argument* result = new IntrinsicFunctor(std::move(fun), Own<Argument>(aggrs[0]), Own<Argument>(rhs));
  
      return result;
  }
  
  /**
   * Returns a vector of arguments that should replace the given argument after one step of inlining.
   * Note: This function is currently generalised to perform any required inlining within aggregators
   * as well, making it simple to extend to this later on if desired (and the semantic check is removed).
   */
  // TODO (azreika): rewrite this method, removing aggregators
  NullableVector<Argument*> getInlinedArgument(Program& program, const Argument* arg) {
      bool changed = false;
      std::vector<Argument*> versions;
  
      // Each argument has to be handled differently - essentially, want to go down to
      // nested aggregators, and inline their bodies if needed.
      if (const auto* aggr = as<Aggregator>(arg)) {
          // First try inlining the target expression if necessary
          if (aggr->getTargetExpression() != nullptr) {
              NullableVector<Argument*> argumentVersions =
                      getInlinedArgument(program, aggr->getTargetExpression());
  
              if (argumentVersions.isValid()) {
                  // An element in the target expression can be inlined!
                  changed = true;
  
                  // Create a new aggregator per version of the target expression
                  for (Argument* newArg : argumentVersions.getVector()) {
                      auto* newAggr = new Aggregator(aggr->getBaseOperator(), Own<Argument>(newArg));
                      VecOwn<Literal> newBody;
                      for (Literal* lit : aggr->getBodyLiterals()) {
                          newBody.push_back(souffle::clone(lit));
                      }
                      newAggr->setBody(std::move(newBody));
                      versions.push_back(newAggr);
                  }
              }
          }
  
          // Try inlining body arguments if the target expression has not been changed.
          // (At this point we only handle one step of inlining at a time)
          if (!changed) {
              std::vector<Literal*> bodyLiterals = aggr->getBodyLiterals();
              for (std::size_t i = 0; i < bodyLiterals.size(); i++) {
                  Literal* currLit = bodyLiterals[i];
  
                  NullableVector<std::vector<Literal*>> literalVersions = getInlinedLiteral(program, currLit);
  
                  if (literalVersions.isValid()) {
                      // Literal can be inlined!
                      changed = true;
  
                      AggregateOp op = aggr->getBaseOperator();
  
                      // Create an aggregator (with the same operation) for each possible body
                      std::vector<Aggregator*> aggrVersions;
                      for (std::vector<Literal*> inlineVersions : literalVersions.getVector()) {
                          Own<Argument> target;
                          if (aggr->getTargetExpression() != nullptr) {
                              target = souffle::clone(aggr->getTargetExpression());
                          }
                          auto* newAggr = new Aggregator(aggr->getBaseOperator(), std::move(target));
  
                          VecOwn<Literal> newBody;
                          // Add in everything except the current literal being replaced
                          for (std::size_t j = 0; j < bodyLiterals.size(); j++) {
                              if (i != j) {
                                  newBody.push_back(souffle::clone(bodyLiterals[j]));
                              }
                          }
  
                          // Add in everything new that replaces that literal
                          for (Literal* addedLit : inlineVersions) {
                              newBody.push_back(Own<Literal>(addedLit));
                          }
  
                          newAggr->setBody(std::move(newBody));
                          aggrVersions.push_back(newAggr);
                      }
  
                      // Utility lambda: get functor used to tie aggregators together.
                      auto aggregateToFunctor = [](AggregateOp op) {
                          switch (op) {
                              case AggregateOp::MIN:
                              case AggregateOp::FMIN:
                              case AggregateOp::UMIN: return "min";
                              case AggregateOp::MAX:
                              case AggregateOp::FMAX:
                              case AggregateOp::UMAX: return "max";
                              case AggregateOp::SUM:
                              case AggregateOp::FSUM:
                              case AggregateOp::USUM:
                              case AggregateOp::COUNT: return "+";
                              case AggregateOp::MEAN: fatal("no translation");
                          }
  
                          UNREACHABLE_BAD_CASE_ANALYSIS
                      };
                      // Create the actual overall aggregator that ties the replacement aggregators together.
                      // example: min x : { a(x) }. <=> min ( min x : { a1(x) }, min x : { a2(x) }, ... )
                      if (op != AggregateOp::MEAN) {
                          versions.push_back(combineAggregators(aggrVersions, aggregateToFunctor(op)));
                      }
                  }
  
                  // Only perform one stage of inlining at a time.
                  if (changed) {
                      break;
                  }
              }
          }
      } else if (const auto* functor = as<Functor>(arg)) {
          std::size_t i = 0;
          for (auto funArg : functor->getArguments()) {
              // TODO (azreika): use unique pointers
              // try inlining each argument from left to right
              NullableVector<Argument*> argumentVersions = getInlinedArgument(program, funArg);
              if (argumentVersions.isValid()) {
                  changed = true;
                  for (Argument* newArgVersion : argumentVersions.getVector()) {
                      // same functor but with new argument version
                      VecOwn<Argument> argsCopy;
                      std::size_t j = 0;
                      for (auto& functorArg : functor->getArguments()) {
                          if (j == i) {
                              argsCopy.emplace_back(newArgVersion);
                          } else {
                              argsCopy.emplace_back(souffle::clone(functorArg));
                          }
                          ++j;
                      }
                      if (const auto* intrFunc = as<IntrinsicFunctor>(arg)) {
                          auto* newFunctor = new IntrinsicFunctor(
                                  intrFunc->getBaseFunctionOp(), std::move(argsCopy), functor->getSrcLoc());
                          versions.push_back(newFunctor);
                      } else if (const auto* userFunc = as<UserDefinedFunctor>(arg)) {
                          auto* newFunctor = new UserDefinedFunctor(
                                  userFunc->getName(), std::move(argsCopy), userFunc->getSrcLoc());
                          versions.push_back(newFunctor);
                      }
                  }
                  // only one step at a time
                  break;
              }
              ++i;
          }
      } else if (const auto* cast = as<ast::TypeCast>(arg)) {
          NullableVector<Argument*> argumentVersions = getInlinedArgument(program, cast->getValue());
          if (argumentVersions.isValid()) {
              changed = true;
              for (Argument* newArg : argumentVersions.getVector()) {
                  Argument* newTypeCast = new ast::TypeCast(Own<Argument>(newArg), cast->getType());
                  versions.push_back(newTypeCast);
              }
          }
      } else if (const auto* record = as<RecordInit>(arg)) {
          std::vector<Argument*> recordArguments = record->getArguments();
          for (std::size_t i = 0; i < recordArguments.size(); i++) {
              Argument* currentRecArg = recordArguments[i];
              NullableVector<Argument*> argumentVersions = getInlinedArgument(program, currentRecArg);
              if (argumentVersions.isValid()) {
                  changed = true;
                  for (Argument* newArgumentVersion : argumentVersions.getVector()) {
                      auto* newRecordArg = new RecordInit();
                      for (std::size_t j = 0; j < recordArguments.size(); j++) {
                          if (i == j) {
                              newRecordArg->addArgument(Own<Argument>(newArgumentVersion));
                          } else {
                              newRecordArg->addArgument(souffle::clone(recordArguments[j]));
                          }
                      }
                      versions.push_back(newRecordArg);
                  }
              }
  
              // Only perform one stage of inlining at a time.
              if (changed) {
                  break;
              }
          }
      }
  
      if (changed) {
          // Return a valid vector - replacements need to be made!
          return NullableVector<Argument*>(versions);
      } else {
          // Return an invalid vector - no inlining has occurred
          return NullableVector<Argument*>();
      }
  }
  
  /**
   * Returns a vector of atoms that should replace the given atom after one step of inlining.
   * Assumes the relation the atom belongs to is not inlined itself.
   */
  NullableVector<Atom*> getInlinedAtom(Program& program, Atom& atom) {
      bool changed = false;
      std::vector<Atom*> versions;
  
      // Try to inline each of the atom's arguments
      std::vector<Argument*> arguments = atom.getArguments();
      for (std::size_t i = 0; i < arguments.size(); i++) {
          Argument* arg = arguments[i];
  
          NullableVector<Argument*> argumentVersions = getInlinedArgument(program, arg);
  
          if (argumentVersions.isValid()) {
              // Argument has replacements
              changed = true;
  
              // Create a new atom per new version of the argument
              for (Argument* newArgument : argumentVersions.getVector()) {
                  auto args = atom.getArguments();
                  VecOwn<Argument> newArgs;
                  for (std::size_t j = 0; j < args.size(); j++) {
                      if (j == i) {
                          newArgs.emplace_back(newArgument);
                      } else {
                          newArgs.emplace_back(souffle::clone(args[j]));
                      }
                  }
                  auto* newAtom = new Atom(atom.getQualifiedName(), std::move(newArgs), atom.getSrcLoc());
                  versions.push_back(newAtom);
              }
          }
  
          // Only perform one stage of inlining at a time.
          if (changed) {
              break;
          }
      }
  
      if (changed) {
          // Return a valid vector - replacements need to be made!
          return NullableVector<Atom*>(versions);
      } else {
          // Return an invalid vector - no replacements need to be made
          return NullableVector<Atom*>();
      }
  }
  
  /**
   * Tries to perform a single step of inlining on the given literal.
   * Returns a pair of nullable vectors (v, w) such that:
   *    - v is valid if and only if the literal can be directly inlined, whereby it
   *      contains the bodies that replace it
   *    - if v is not valid, then w is valid if and only if the literal cannot be inlined directly,
   *      but contains a subargument that can be. In this case, it will contain the versions
   *      that will replace it.
   *    - If both are invalid, then no more inlining can occur on this literal and we are done.
   */
  NullableVector<std::vector<Literal*>> getInlinedLiteral(Program& program, Literal* lit) {
      bool inlined = false;
      bool changed = false;
  
      std::vector<std::vector<Literal*>> addedBodyLiterals;
      std::vector<Literal*> versions;
  
      if (auto* atom = as<Atom>(lit)) {
          // Check if this atom is meant to be inlined
          Relation* rel = getRelation(program, atom->getQualifiedName());
  
          if (rel->hasQualifier(RelationQualifier::INLINE)) {
              // We found an atom in the clause that needs to be inlined!
              // The clause needs to be replaced
              inlined = true;
  
              // N new clauses should be formed, where N is the number of clauses
              // associated with the inlined relation
              for (Clause* inClause : getClauses(program, *rel)) {
                  // Form the replacement clause
                  std::pair<NullableVector<Literal*>, std::vector<BinaryConstraint*>> inlineResult =
                          inlineBodyLiterals(atom, inClause);
                  NullableVector<Literal*> replacementBodyLiterals = inlineResult.first;
                  std::vector<BinaryConstraint*> currConstraints = inlineResult.second;
  
                  if (!replacementBodyLiterals.isValid()) {
                      // Failed to unify the atoms! We can skip this one...
                      continue;
                  }
  
                  // Unification successful - the returned vector of literals represents one possible body
                  // replacement We can add in the unification constraints as part of these literals.
                  std::vector<Literal*> bodyResult = replacementBodyLiterals.getVector();
  
                  for (BinaryConstraint* cons : currConstraints) {
                      bodyResult.push_back(cons);
                  }
  
                  addedBodyLiterals.push_back(bodyResult);
              }
          } else {
              // Not meant to be inlined, but a subargument may be
              NullableVector<Atom*> atomVersions = getInlinedAtom(program, *atom);
              if (atomVersions.isValid()) {
                  // Subnode needs to be inlined, so we have a vector of replacement atoms
                  changed = true;
                  for (Atom* newAtom : atomVersions.getVector()) {
                      versions.push_back(newAtom);
                  }
              }
          }
      } else if (auto neg = as<Negation>(lit)) {
          // For negations, check the corresponding atom
          Atom* atom = neg->getAtom();
          NullableVector<std::vector<Literal*>> atomVersions = getInlinedLiteral(program, atom);
  
          if (atomVersions.isValid()) {
              // The atom can be inlined
              inlined = true;
  
              if (atomVersions.getVector().empty()) {
                  // No clauses associated with the atom, so just becomes a true literal
                  addedBodyLiterals.push_back({new BooleanConstraint(true)});
              } else {
                  // Suppose an atom a(x) is inlined and has the following rules:
                  //  - a(x) :- a11(x), a12(x).
                  //  - a(x) :- a21(x), a22(x).
                  // Then, a(x) <- (a11(x) ^ a12(x)) v (a21(x) ^ a22(x))
                  //  => !a(x) <- (!a11(x) v !a12(x)) ^ (!a21(x) v !a22(x))
                  //  => !a(x) <- (!a11(x) ^ !a21(x)) v (!a11(x) ^ !a22(x)) v ...
                  // Essentially, produce every combination (m_1 ^ m_2 ^ ...) where m_i is a
                  // negated literal in the ith rule of a.
                  addedBodyLiterals = formNegatedLiterals(program, atom);
              }
          }
          if (atomVersions.isValid()) {
              for (const auto& curVec : atomVersions.getVector()) {
                  for (auto* cur : curVec) {
                      delete cur;
                  }
              }
          }
      } else if (auto* constraint = as<BinaryConstraint>(lit)) {
          NullableVector<Argument*> lhsVersions = getInlinedArgument(program, constraint->getLHS());
          if (lhsVersions.isValid()) {
              changed = true;
              for (Argument* newLhs : lhsVersions.getVector()) {
                  Literal* newLit = new BinaryConstraint(constraint->getBaseOperator(), Own<Argument>(newLhs),
                          souffle::clone(constraint->getRHS()));
                  versions.push_back(newLit);
              }
          } else {
              NullableVector<Argument*> rhsVersions = getInlinedArgument(program, constraint->getRHS());
              if (rhsVersions.isValid()) {
                  changed = true;
                  for (Argument* newRhs : rhsVersions.getVector()) {
                      Literal* newLit = new BinaryConstraint(constraint->getBaseOperator(),
                              souffle::clone(constraint->getLHS()), Own<Argument>(newRhs));
                      versions.push_back(newLit);
                  }
              }
          }
      }
  
      if (changed) {
          // Not inlined directly but found replacement literals
          // Rewrite these as single-literal bodies
          for (Literal* version : versions) {
              std::vector<Literal*> newBody;
              newBody.push_back(version);
              addedBodyLiterals.push_back(newBody);
          }
          inlined = true;
      }
  
      if (inlined) {
          return NullableVector<std::vector<Literal*>>(addedBodyLiterals);
      } else {
          return NullableVector<std::vector<Literal*>>();
      }
  }
  
  /**
   * Returns a list of clauses that should replace the given clause after one step of inlining.
   * If no inlining can occur, the list will only contain a clone of the original clause.
   */
  std::vector<Clause*> getInlinedClause(Program& program, const Clause& clause) {
      bool changed = false;
      std::vector<Clause*> versions;
  
      // Try to inline things contained in the arguments of the head first.
      // E.g. `a(x, max y : { b(y) }) :- c(x).`, where b should be inlined.
      Atom* head = clause.getHead();
      NullableVector<Atom*> headVersions = getInlinedAtom(program, *head);
      if (headVersions.isValid()) {
          // The head atom can be inlined!
          changed = true;
  
          // Produce the new clauses with the replacement head atoms
          for (Atom* newHead : headVersions.getVector()) {
              auto newClause = mk<Clause>(Own<Atom>(newHead), clause.getSrcLoc());
              newClause->setBodyLiterals(souffle::clone(clause.getBodyLiterals()));
  
              // FIXME: tomp - hack - this should be managed
              versions.push_back(newClause.release());
          }
      }
  
      // Only perform one stage of inlining at a time.
      // If the head atoms did not need inlining, try inlining atoms nested in the body.
      if (!changed) {
          std::vector<Literal*> bodyLiterals = clause.getBodyLiterals();
          for (std::size_t i = 0; i < bodyLiterals.size(); i++) {
              Literal* currLit = bodyLiterals[i];
  
              // Three possible cases when trying to inline a literal:
              //  1) The literal itself may be directly inlined. In this case, the atom can be replaced
              //    with multiple different bodies, as the inlined atom may have several rules.
              //  2) Otherwise, the literal itself may not need to be inlined, but a subnode (e.g. an argument)
              //    may need to be inlined. In this case, an altered literal must replace the original.
              //    Again, several possible versions may exist, as the inlined relation may have several rules.
              //  3) The literal does not depend on any inlined relations, and so does not need to be changed.
              NullableVector<std::vector<Literal*>> litVersions = getInlinedLiteral(program, currLit);
  
              if (litVersions.isValid()) {
                  // Case 1 and 2: Inlining has occurred!
                  changed = true;
  
                  // The literal may be replaced with several different bodies.
                  // Create a new clause for each possible version.
                  std::vector<std::vector<Literal*>> bodyVersions = litVersions.getVector();
  
                  // Create the base clause with the current literal removed
                  auto baseClause = cloneHead(clause);
                  for (Literal* oldLit : bodyLiterals) {
                      if (currLit != oldLit) {
                          baseClause->addToBody(souffle::clone(oldLit));
                      }
                  }
  
                  for (std::vector<Literal*> const& body : bodyVersions) {
                      auto replacementClause = souffle::clone(baseClause);
  
                      // Add in the current set of literals replacing the inlined literal
                      // In Case 2, each body contains exactly one literal
                      replacementClause->addToBody(VecOwn<Literal>(body.begin(), body.end()));
  
                      // FIXME: This is a horrible hack.  Should convert
                      // versions to hold Own<>
                      versions.push_back(replacementClause.release());
                  }
              }
  
              // Only replace at most one literal per iteration
              if (changed) {
                  break;
              }
          }
      }
  
      if (!changed) {
          // Case 3: No inlining changes, so a clone of the original should be returned
          // FIXME: This is a horrible hack.  Should convert
          // res to hold Own<>
          std::vector<Clause*> ret;
          ret.push_back(souffle::clone(clause).release());
          return ret;
      } else {
          // Inlining changes, so return the replacement clauses.
          return versions;
      }
  }
  
  bool InlineRelationsTransformer::transform(TranslationUnit& translationUnit) {
      bool changed = false;
      Program& program = translationUnit.getProgram();
  
      // Replace constants in the head of inlined clauses with (constrained) variables.
      // This is done to simplify atom unification, particularly when negations are involved.
      changed |= normaliseInlinedHeads(program);
  
      // Remove underscores in inlined atoms in the program to avoid issues during atom unification
      changed |= nameInlinedUnderscores(program);
  
      // Keep trying to inline things until we reach a fixed point.
      // Since we know there are no cyclic dependencies between inlined relations, this will necessarily
      // terminate.
      bool clausesChanged = true;
      while (clausesChanged) {
          std::set<Clause*> clausesToDelete;
          clausesChanged = false;
  
          // Go through each relation in the program and check if we need to inline any of its clauses
          for (Relation* rel : program.getRelations()) {
              // Skip if the relation is going to be inlined
              if (rel->hasQualifier(RelationQualifier::INLINE)) {
                  continue;
              }
  
              // Go through the relation's clauses and try inlining them
              for (Clause* clause : getClauses(program, *rel)) {
                  if (containsInlinedAtom(program, *clause)) {
                      // Generate the inlined versions of this clause - the clause will be replaced by these
                      std::vector<Clause*> newClauses = getInlinedClause(program, *clause);
  
                      // Replace the clause with these equivalent versions
                      clausesToDelete.insert(clause);
                      for (Clause* replacementClause : newClauses) {
                          program.addClause(Own<Clause>(replacementClause));
                      }
  
                      // We've changed the program this iteration
                      clausesChanged = true;
                      changed = true;
                  }
              }
          }
  
          // Delete all clauses that were replaced
          for (const Clause* clause : clausesToDelete) {
              program.removeClause(clause);
              changed = true;
          }
      }
  
      return changed;
  }
  
  }  // namespace souffle::ast::transform
  