Partition.h

Go to the documentation of this file.

#pragma once
 
#include <cstddef>
#include <map>
#include <ranges>
#include <set>
#include <span>
 
#include "storm/storage/BitVector.h"
#include "storm/utility/macros.h"
 
namespace storm::bisimulation {

class Partition {
   public:
    using ElementIndex = uint64_t;

    using Block = std::span<ElementIndex const>;

    struct BlockCompare {
        bool operator()(Block const& lhs, Block const& rhs) const {
            if (lhs.size() < rhs.size()) {
                return true;
            }
            if (lhs.size() > rhs.size()) {
                return false;
            }
            return lhs.data() < rhs.data();
        }
    };

    using OrderedBlockSet = std::set<Block, BlockCompare>;

    template<typename T>
    using OrderedBlockMap = std::map<Block, T, BlockCompare>;

    class NonSuperBlockSet {
       public:
        explicit NonSuperBlockSet(Partition const& partition);

        std::size_t size() const;

        bool empty() const;

        bool contains(Block const& block) const;

        void insert(Block const& block);

        Block pop();
 
       private:
        Partition const& partition;
        std::vector<uint64_t> containedBlockIndices;
        storm::storage::BitVector blockIndices;
    };
 
    Partition() = default;
    Partition(Partition const& other) = default;
    Partition& operator=(Partition const& other) = default;
    Partition(Partition&& other) noexcept = default;
    Partition& operator=(Partition&& other) = default;

    explicit Partition(ElementIndex numElements);

    std::size_t getNumberOfBlocks() const;

    std::size_t getNumberOfElements() const;

    Block getUniversalBlock() const;

    Block getBlockOfElement(ElementIndex element) const;

    bool contains(Block const& block, ElementIndex element) const;

    bool isSubBlockOf(Block const& subblock, Block const& superblock) const;

    bool isEqualBlock(Block const& lhs, Block const& rhs) const;

    bool isSameBlock(ElementIndex const& element1, ElementIndex const& element2) const;

    bool isBlockOfElement(Block const& block, ElementIndex const& element) const;

    bool checkBlockValidity(Block const& block) const;

    bool isProperSuperBlock(Block const& block) const;

    template<typename Func>
        requires std::invocable<Func, Block>
    void forEachBlock(Func const& f) const {
        forEachSubBlock(getUniversalBlock(), f);
    }

    template<typename Func>
        requires std::invocable<Func, Block>
    void forEachSubBlock(Block const& superBlock, Func const& f) const {
        STORM_LOG_ASSERT(!superBlock.empty(), "Superblock is empty.");
        auto superBlockEnd = getBlockIndex(superBlock) + superBlock.size();
        for (BlockIndex blockStart = getBlockIndex(superBlock); blockStart < superBlockEnd;) {
            auto const& subBlock = getBlockOfElement(blockContents[blockStart]);
            blockStart += subBlock.size();
            f(subBlock);
        }
    }

    template<typename SplittingOrder>
        requires std::invocable<SplittingOrder, ElementIndex, ElementIndex>
    bool splitBlockByOrder(Block const& block, SplittingOrder const& less) {
        return splitBlockByOrder(block, less, less);
    }

    template<typename Order, typename Condition>
        requires std::invocable<Order, ElementIndex, ElementIndex> && std::invocable<Condition, ElementIndex, ElementIndex>
    bool splitBlockByOrder(Block const& block, Order const& less, Condition const& splitCondition) {
        if (block.size() <= 1) {
            return false;  // nothing to do
        }
 
        // Sort the contents of the block first
        STORM_LOG_ASSERT(!isProperSuperBlock(block), "Tried to split a block that consists of multiple sub-blocks.");
        auto const blockStart = getBlockIndex(block);
 
        auto const blockEnd = blockStart + block.size();
        std::sort(blockContents.begin() + blockStart, blockContents.begin() + blockEnd, less);
 
        // update the inverse after sorting
        for (ElementIndex i = blockStart; i < blockEnd; ++i) {
            blockContentsInverse[blockContents[i]] = i;
        }
 
        // Catch the special case where there is no split
        if (!splitCondition(blockContents[blockStart], blockContents[blockEnd - 1])) {
            return false;  // nothing to do
        }
 
        // helper function to find the end index of a current block
        auto getEndOfBlock = [this, splitCondition, blockEnd](BlockIndex const currIndex) {
            for (auto i = currIndex + 1; i < blockEnd; ++i) {
                if (splitCondition(blockContents[currIndex], blockContents[i])) {
                    return i;
                }
            }
            return blockEnd;
        };
 
        // Now create new blocks whenever the order changes
        for (auto newBlockIndex = blockStart; newBlockIndex < blockEnd;) {
            auto const newBlockEnd = getEndOfBlock(newBlockIndex);
            registerNewBlock(newBlockIndex, newBlockEnd);
            newBlockIndex = newBlockEnd;
            ++numBlocks;
        }
        --numBlocks;  // the old superblock is no longer a block as counted by this partition
        STORM_LOG_ASSERT(isProperSuperBlock(block), "Partition in inconsistent state: Block was not split into multiple sub-blocks.");
        STORM_LOG_ASSERT(checkBlockValidity(block), "Partition in inconsistent state: Block is not valid.");
        return true;  // there must have been a split because the case without a split is already caught above
    }

    template<typename Condition>
        requires std::invocable<Condition, ElementIndex, ElementIndex>
    bool splitBlockByClustering(Block const& block, Condition const& condition) {
        STORM_LOG_ASSERT(!isProperSuperBlock(block), "Tried to split a block that consists of multiple sub-blocks.");
        if (block.size() <= 1) {
            return false;  // nothing to do
        }
 
        auto const blockStart = getBlockIndex(block);
        auto const blockEnd = blockStart + block.size();
 
        // Swaps the elements in [start, blockEnd) so that those compatible (!condition) with the anchor at start directly follow it.
        // @return mid such that all elements in [start, mid) are compatible with the anchor (located at start) and those at [mid, blockEnd) are not.
        auto const partitionBlock = [this, &condition, &blockEnd](BlockIndex const start) {
            // The anchor itself stays in place, so we only partition the elements after it.
            auto const anchor = blockContents[start];
            return partitionBlockContents(start + 1, blockEnd, [&condition, &anchor](ElementIndex const element) { return condition(anchor, element); });
        };
 
        for (uint64_t start = blockStart; start < blockEnd;) {
            auto const currEnd = partitionBlock(start);
            // Catch the special case where there is no split
            if (start == blockStart && currEnd == blockEnd) {
                return false;  // nothing to do
            }
            registerNewBlock(start, currEnd);
            start = currEnd;
            ++numBlocks;
        }
        --numBlocks;  // the old superblock is no longer a block as counted by this partition
        STORM_LOG_ASSERT(isProperSuperBlock(block), "Partition in inconsistent state: Block was not split into multiple sub-blocks.");
        STORM_LOG_ASSERT(checkBlockValidity(block), "Partition in inconsistent state: Block is not valid.");
        return true;  // there must have been a split because the case without a split is already caught above
    }

    template<typename SplittingPredicate>
        requires std::predicate<SplittingPredicate, ElementIndex>
    std::pair<Block, Block> splitBlockByPredicate(Block const& block, SplittingPredicate const& f) {
        STORM_LOG_ASSERT(!isProperSuperBlock(block), "Tried to split a block that consists of multiple sub-blocks.");
        STORM_LOG_ASSERT(!block.empty(), "Tried to split an empty block");
 
        // swap the block contents so that all elements e with f(e) == false come first
        auto const blockStart = getBlockIndex(block);
        auto const blockEnd = blockStart + block.size();
        auto const l = partitionBlockContents(blockStart, blockEnd, f);
 
        // Handle cases where there is no split
        if (l == blockStart) {
            return {Block{}, block};  // all elements are true
        } else if (l == blockEnd) {
            return {block, Block{}};  // all elements are false
        }
 
        // Perform a split
        Block noBlock = registerNewBlock(blockStart, l);
        Block yesBlock = registerNewBlock(l, blockEnd);
        ++numBlocks;  // We have split a single block into two blocks, so the total number increments by 1
        return {noBlock, yesBlock};
    }

    template<typename SplitRange>
        requires std::ranges::input_range<std::remove_reference_t<SplitRange>> &&
                 std::same_as<std::ranges::range_value_t<std::remove_reference_t<SplitRange>>, ElementIndex>
    std::pair<Block, Block> splitBlockByRange(Block const& block, SplitRange&& r) {
        STORM_LOG_ASSERT(!isProperSuperBlock(block), "Tried to split a block that consists of multiple sub-blocks.");
        STORM_LOG_ASSERT(!block.empty(), "Tried to split an empty block");
 
        auto const blockStart = getBlockIndex(block);
        auto const blockEnd = blockStart + block.size();
 
        // We split the block into a "no"-part and a "yes"-part
 
        auto yesBlockStart = blockEnd;
        for (ElementIndex const element : r) {
            auto const src = blockContentsInverse[element];
            if (src < blockStart || src >= yesBlockStart) {
                continue;  // element is either not in the block or already occurred before (duplicate in r)
            }
            // swap the element to the end of the block, into the 'yes' part
            --yesBlockStart;
            std::swap(blockContents[src], blockContents[yesBlockStart]);
            blockContentsInverse[blockContents[src]] = src;
            blockContentsInverse[blockContents[yesBlockStart]] = yesBlockStart;
        }
 
        // Handle cases where there is no split
        if (yesBlockStart == blockStart) {
            return {Block{}, block};  // all elements are in the range
        } else if (yesBlockStart == blockEnd) {
            return {block, Block{}};  // no elements are in the range
        }
 
        // Perform a split
        Block noBlock = registerNewBlock(blockStart, yesBlockStart);
        Block yesBlock = registerNewBlock(yesBlockStart, blockEnd);
        ++numBlocks;  // We have split a single block into two blocks, so the total number increments by 1
        return {noBlock, yesBlock};
    }
 
   private:
    using BlockIndex = uint64_t;

    template<typename SplittingPredicate>
        requires std::predicate<SplittingPredicate, ElementIndex>
    BlockIndex partitionBlockContents(BlockIndex const start, BlockIndex const end, SplittingPredicate const& f) {
        STORM_LOG_ASSERT(start <= end, "Tried to partition an invalid range.");
        if (start == end) {
            return start;  // empty range, nothing to swap
        }
        auto l = start;
        auto r = end - 1;
        // Loop invariant: all elements at position < l are false, all elements at position > r are true
        while (l <= r) {
            while (l <= r && !f(std::as_const(blockContents[l]))) {
                ++l;
            }
            if (l > r) {
                break;  // no more elements to swap
            }
            // At this point we know that f(blockContents[l]) == true
            while (l < r && f(std::as_const(blockContents[r]))) {
                --r;
            }
            if (l == r) {
                // We have f(blockContents[r]) == f(blockContents[l]) == true
                --r;
                break;  // l > r holds now
            }
            std::swap(blockContents[l], blockContents[r]);
            blockContentsInverse[blockContents[l]] = l;
            blockContentsInverse[blockContents[r]] = r;
            ++l;
            --r;
        }
        STORM_LOG_ASSERT(l == r + 1, "Unexpected indices");
        return l;
    }

    BlockIndex getBlockIndex(Block const& block) const;

    Block getBlockFromIndex(BlockIndex blockIndex) const;

    Block registerNewBlock(BlockIndex const start, BlockIndex const end);

    std::vector<ElementIndex> blockContents;

    std::vector<BlockIndex> blockContentsInverse;

    std::vector<BlockIndex> elementToBlockIndex;

    std::vector<std::size_t> blockSizes;

    uint64_t numBlocks = 0;
};
 
std::ostream& operator<<(std::ostream& os, const Partition& partition);
 
}  // namespace storm::bisimulation