Aggregate.cpp

Go to the documentation of this file.

/*
 * Souffle - A Datalog Compiler
 * Copyright (c) 2020, 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 Ground.cpp
 *
 * Implements Aggregate Analysis methods to determine scope of variables
 *
 ***********************************************************************/
 
#include "ast/analysis/Aggregate.h"
#include "ast/Aggregator.h"
#include "ast/Atom.h"
#include "ast/Clause.h"
#include "ast/Relation.h"
#include "ast/TranslationUnit.h"
#include "ast/Variable.h"
#include "ast/analysis/Ground.h"
#include "ast/utility/NodeMapper.h"
#include "ast/utility/Utils.h"
#include "ast/utility/Visitor.h"
#include "souffle/utility/StringUtil.h"
 
#include <algorithm>
#include <map>
#include <memory>
#include <ostream>
#include <set>
#include <string>
#include <utility>
#include <vector>
 
namespace souffle::ast::analysis {
 
/**
 *  Computes the set of local variables
 *  in an aggregate expression.
 *  This is just the set of all variables occurring in the aggregate
 *  MINUS the injected variables
 *  MINUS the witness variables. :)
 *
 **/
std::set<std::string> getLocalVariables(
        const TranslationUnit& tu, const Clause& clause, const Aggregator& aggregate) {
    std::set<std::string> allVariablesInAggregate;
    visitDepthFirst(aggregate, [&](const Variable& v) { allVariablesInAggregate.insert(v.getName()); });
    std::set<std::string> injectedVariables = getInjectedVariables(tu, clause, aggregate);
    std::set<std::string> witnessVariables = getWitnessVariables(tu, clause, aggregate);
    std::set<std::string> localVariables;
    for (const std::string& name : allVariablesInAggregate) {
        if (injectedVariables.find(name) == injectedVariables.end() &&
                witnessVariables.find(name) == witnessVariables.end()) {
            localVariables.insert(name);
        }
    }
    return localVariables;
}
/**
 * Computes the set of witness variables that are used in the aggregate
 * A variable is a witness if it occurs in the aggregate body (but not in an inner aggregate)
 * and also occurs ungrounded in the outer scope.
 **/
std::set<std::string> getWitnessVariables(
        const TranslationUnit& tu, const Clause& clause, const Aggregator& aggregate) {
    // 1. Create an aggregatorless clause so that we can analyse the groundings
    // in the rule scope. This is one where we replace every aggregator with
    // a variable, and then provide a grounding for that variable.
    struct M : public NodeMapper {
        // Variables introduced to replace aggregators
        mutable std::set<std::string> aggregatorVariables;
 
        const std::set<std::string>& getAggregatorVariables() {
            return aggregatorVariables;
        }
 
        std::unique_ptr<Node> operator()(std::unique_ptr<Node> node) const override {
            static int numReplaced = 0;
            if (dynamic_cast<Aggregator*>(node.get()) != nullptr) {
                // Replace the aggregator with a variable
                std::stringstream newVariableName;
                newVariableName << "+aggr_var_" << numReplaced++;
 
                // Keep track of which variables are bound to aggregators
                aggregatorVariables.insert(newVariableName.str());
 
                return std::make_unique<Variable>(newVariableName.str());
            }
            node->apply(*this);
            return node;
        }
    };
 
    auto aggregatorlessClause = std::make_unique<Clause>();
    aggregatorlessClause->setHead(std::make_unique<Atom>("*"));
    for (Literal* lit : clause.getBodyLiterals()) {
        aggregatorlessClause->addToBody(souffle::clone(lit));
    }
 
    auto negatedHead = std::make_unique<Negation>(souffle::clone(clause.getHead()));
    aggregatorlessClause->addToBody(std::move(negatedHead));
 
    // Replace all aggregates with variables
    M update;
    aggregatorlessClause->apply(update);
    auto groundingAtom = std::make_unique<Atom>("+grounding_atom");
    for (std::string variableName : update.getAggregatorVariables()) {
        groundingAtom->addArgument(std::make_unique<Variable>(variableName));
    }
    aggregatorlessClause->addToBody(std::move(groundingAtom));
    // 2. Create an aggregate clause so that we can check
    // that it IS this aggregate giving a grounding to the candidate variable.
    auto aggregateSubclause = std::make_unique<Clause>();
    aggregateSubclause->setHead(mk<Atom>("*"));
    for (const auto& lit : aggregate.getBodyLiterals()) {
        aggregateSubclause->addToBody(souffle::clone(lit));
    }
 
    std::set<std::string> witnessVariables;
    auto isGroundedInAggregateSubclause = analysis::getGroundedTerms(tu, *aggregateSubclause);
    // 3. Calculate all the witness variables
    // A witness will occur ungrounded in the aggregatorlessClause
    for (const auto& argPair : analysis::getGroundedTerms(tu, *aggregatorlessClause)) {
        if (const auto* variable = dynamic_cast<const Variable*>(argPair.first)) {
            bool variableIsGrounded = argPair.second;
            if (!variableIsGrounded) {
                // then we expect it to be grounded in the aggregate subclause
                // if it's a witness!!
                for (const auto& aggArgPair : isGroundedInAggregateSubclause) {
                    if (const auto* var = dynamic_cast<const Variable*>(aggArgPair.first)) {
                        bool aggVariableIsGrounded = aggArgPair.second;
                        if (var->getName() == variable->getName() && aggVariableIsGrounded) {
                            witnessVariables.insert(variable->getName());
                        }
                    }
                }
            }
        }
    }
    // 4. A witness variable may actually "originate" from an outer scope and
    // just have been injected into this inner aggregate. Just check the set of injected variables
    // and quickly minus them out.
    std::set<std::string> injectedVariables = analysis::getInjectedVariables(tu, clause, aggregate);
    for (const std::string& injected : injectedVariables) {
        witnessVariables.erase(injected);
    }
 
    return witnessVariables;
}
/**
 *  Computes the set of variables occurring outside the aggregate
 **/
std::set<std::string> getVariablesOutsideAggregate(const Clause& clause, const Aggregator& aggregate) {
    std::map<std::string, int> variableOccurrences;
    visitDepthFirst(clause, [&](const Variable& var) { variableOccurrences[var.getName()]++; });
    visitDepthFirst(aggregate, [&](const Variable& var) { variableOccurrences[var.getName()]--; });
    std::set<std::string> variablesOutsideAggregate;
    for (auto const& pair : variableOccurrences) {
        std::string v = pair.first;
        int numOccurrences = pair.second;
        if (numOccurrences > 0) {
            variablesOutsideAggregate.insert(v);
        }
    }
    return variablesOutsideAggregate;
}
 
std::string findUniqueVariableName(const Clause& clause, std::string base) {
    std::set<std::string> variablesInClause;
    visitDepthFirst(clause, [&](const Variable& v) { variablesInClause.insert(v.getName()); });
    int varNum = 0;
    std::string candidate = base;
    while (variablesInClause.find(candidate) != variablesInClause.end()) {
        candidate = base + toString(varNum++);
    }
    return candidate;
}
 
std::string findUniqueRelationName(const Program& program, std::string base) {
    int counter = 0;
    auto candidate = base;
    while (getRelation(program, candidate) != nullptr) {
        candidate = base + toString(counter++);
    }
    return candidate;
}
 
/**
 *  Given an aggregate and a clause, we find all the variables that have been
 *  injected into the aggregate.
 *  This means that the variable occurs grounded in an outer scope.
 *  BUT does not occur in the target expression.
 **/
std::set<std::string> getInjectedVariables(
        const TranslationUnit& tu, const Clause& clause, const Aggregator& aggregate) {
    /**
     *  Imagine we have this:
     *
     *   A(letter, x) :-
     *       a(letter, _),
     *       x = min y : { a(l, y),
     *       y < max y : { a(letter, y) } },
     *       z = max y : {B(y) }.
     *
     *   We need to figure out the base literal in which this aggregate occurs,
     *   so that we don't delete them
     *   Other aggregates occurring in the rule are irrelevant, so we'll replace
     *   the aggregate with a dummy aggregate variable. Then after that,
     *
     *   2. A(letter, x) :-
     *       a(letter, _),
     *       x = min y : { a(l, y),
     *       y < max y : { a(letter, y) } },
     *       z = aggr_var_0.
     *
     *   We should also ground all the aggr_vars so that this grounding
     *   transfers to whatever variable they might be assigned to.
     *
     *   A(letter, x) :-
     *       a(letter, _),
     *       x = min y : { a(l, y),
     *       y < max y : { a(letter, y) } },
     *       z = aggr_var_0,
     *       grounding_atom(aggr_var_0).
     *
     *   Remember to negate the head atom and add it to the body. The clause we deal with will have a
     *   dummy head (*())
     *
     *    *() :- !A(letter, x), a(letter, _), x = min y : { a(l, y), y < max y : { a(letter, y) } },
     *        z = aggr_var_0, grounding_atom(aggr_var_0).
     *
     *   Replace the original aggregate with the same shit because we can't be counting local variables n shit
     *
     *    *() :- !A(letter, x), a(letter, _), x = min y : { a(l, y), y < aggr_var_1 } }, z = aggr_var_0,
     *        grounding_atom(aggr_var_0).
     *
     *   Now, as long as we recorded all of the variables that occurred inside the aggregate, the intersection
     *   of these two sets (variables in an outer scope) and (variables occurring in t
     **/
 
    /**
     * A variable is considered *injected* into an aggregate if it occurs both within the aggregate AS WELL
     * AS *grounded* in the outer scope. An outer scope could be an outer aggregate, or also just the
     *base/rule scope. (i.e. not inside of another aggregate that is not an ancestor of this one).
     *
     *   So then in order to calculate the set of variables that have been injected into an aggregate, we
     * perform the following steps:
     *
     *   1. Find the set of variables occurring within the aggregate "aggregate"
     *   2. Find the set of variables that occur grounded in the outside scope
     *       2a. Replace all non-ancestral aggregates with variables so that we can ignore their variable
     * contents (it is not relevant). 2b. Replace the target aggregate with a variable as well 2c. Ground all
     * occurrences of the non-ancestral aggregate variable replacements (i.e. grounding_atom(+aggr_var_0)) 2d.
     * Visit all variables occurring in this tweaked clause that we made from steps 2a-2c, and check if they
     * are grounded, and also occur in the set of target aggregate variables. Then it is an injected variable.
     **/
    // Step 1
    std::set<std::string> variablesInTargetAggregate;
    visitDepthFirst(aggregate,
            [&](const Variable& variable) { variablesInTargetAggregate.insert(variable.getName()); });
 
    std::set<Own<Aggregator>> ancestorAggregates;
 
    visitDepthFirst(clause, [&](const Aggregator& ancestor) {
        visitDepthFirst(ancestor, [&](const Aggregator& agg) {
            if (agg == aggregate) {
                ancestorAggregates.insert(souffle::clone(&ancestor));
            }
        });
    });
 
    // 1. Replace non-ancestral aggregates with variables
    struct ReplaceAggregatesWithVariables : public NodeMapper {
        // Variables introduced to replace aggregators
        mutable std::set<std::string> aggregatorVariables;
        // ancestors are never replaced so don't need to clone, but actually
        // it will have a child that will become invalid at one point. This is a concern. Need to clone these
        // bastards too.
        std::set<Own<Aggregator>> ancestors;
        // make sure you clone the agg first.
        Own<Aggregator> targetAggregate;
 
        const std::set<std::string>& getAggregatorVariables() {
            return aggregatorVariables;
        }
 
        ReplaceAggregatesWithVariables(std::set<Own<Aggregator>> ancestors, Own<Aggregator> targetAggregate)
                : ancestors(std::move(ancestors)), targetAggregate(std::move(targetAggregate)) {}
 
        std::unique_ptr<Node> operator()(std::unique_ptr<Node> node) const override {
            static int numReplaced = 0;
            if (auto* aggregate = dynamic_cast<Aggregator*>(node.get())) {
                // If we come across an aggregate that is NOT an ancestor of
                // the target aggregate, or that IS itself the target aggregate,
                // we should replace it with a dummy variable.
                bool isAncestor = false;
                for (auto& ancestor : ancestors) {
                    if (*ancestor == *aggregate) {
                        isAncestor = true;
                        break;
                    }
                }
                if (!isAncestor || *aggregate == *targetAggregate) {
                    // Replace the aggregator with a variable
                    std::stringstream newVariableName;
                    newVariableName << "+aggr_var_" << numReplaced++;
                    // Keep track of which variables are bound to aggregators
                    // (only applicable to non-ancestral aggregates)
                    if (!isAncestor) {
                        // we don't want to ground the target aggregate
                        // (why?) because these variables could be injected
                        // and that is fine.
                        aggregatorVariables.insert(newVariableName.str());
                        // but we are not supposed to be judging the aggregate as grounded and injected into
                        // the aggregate, that wouldn't make sense
                    }
                    return mk<Variable>(newVariableName.str());
                }
            }
            node->apply(*this);
            return node;
        }
    };
    // 2. make a clone of the clause and then apply that mapper onto it
    auto clauseCopy = souffle::clone(&clause);
    auto tweakedClause = mk<Clause>();
    // put a fake head here
    tweakedClause->setHead(mk<Atom>("*"));
    // copy in all the old body literals
    for (Literal* lit : clause.getBodyLiterals()) {
        tweakedClause->addToBody(souffle::clone(lit));
    }
    // copy in the head as a negated atom
    tweakedClause->addToBody(mk<Negation>(souffle::clone(clause.getHead())));
    // copy in body literals and also add the old head as a negated atom
    ReplaceAggregatesWithVariables update(std::move(ancestorAggregates), souffle::clone(&aggregate));
    tweakedClause->apply(update);
    // the update will now tell us which variables we need to ground!
    auto groundingAtom = mk<Atom>("+grounding_atom");
    for (std::string variableName : update.getAggregatorVariables()) {
        groundingAtom->addArgument(mk<Variable>(variableName));
    }
    // add the newly created grounding atom to the body
    tweakedClause->addToBody(std::move(groundingAtom));
 
    std::set<std::string> injectedVariables;
    // Search through the tweakedClause to find groundings!
    for (const auto& argPair : analysis::getGroundedTerms(tu, *tweakedClause)) {
        if (const auto* variable = dynamic_cast<const Variable*>(argPair.first)) {
            bool varIsGrounded = argPair.second;
            if (varIsGrounded && variablesInTargetAggregate.find(variable->getName()) !=
                                         variablesInTargetAggregate.end()) {
                // then we have found an injected variable, lovely!!
                injectedVariables.insert(variable->getName());
            }
        }
    }
    // Remove any variables that occur in the target expression of the aggregate
    if (aggregate.getTargetExpression() != nullptr) {
        visitDepthFirst(*aggregate.getTargetExpression(),
                [&](const Variable& v) { injectedVariables.erase(v.getName()); });
    }
 
    return injectedVariables;
}
 
}  // namespace souffle::ast::analysis