1 //===-- llvm/ADT/CombinationGenerator.h ------------------------*- C++ -*--===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 ///
9 /// \file
10 /// Combination generator.
11 ///
12 /// Example: given input {{0, 1}, {2}, {3, 4}} it will produce the following
13 /// combinations: {0, 2, 3}, {0, 2, 4}, {1, 2, 3}, {1, 2, 4}.
14 ///
15 /// It is useful to think of input as vector-of-vectors, where the
16 /// outer vector is the variable space, and inner vector is choice space.
17 /// The number of choices for each variable can be different.
18 ///
19 /// As for implementation, it is useful to think of this as a weird number,
20 /// where each digit (==variable) may have different base (==number of choices).
21 /// Thus modelling of 'produce next combination' is exactly analogous to the
22 /// incrementing of an number - increment lowest digit (pick next choice for the
23 /// variable), and if it wrapped to the beginning then increment next digit.
24 ///
25 //===----------------------------------------------------------------------===//
26 
27 #ifndef LLVM_ADT_COMBINATIONGENERATOR_H
28 #define LLVM_ADT_COMBINATIONGENERATOR_H
29 
30 #include "llvm/ADT/ArrayRef.h"
31 #include "llvm/ADT/STLFunctionalExtras.h"
32 #include "llvm/ADT/SmallVector.h"
33 #include <cassert>
34 #include <cstring>
35 
36 namespace llvm {
37 
38 template <typename choice_type, typename choices_storage_type,
39           int variable_smallsize>
40 class CombinationGenerator {
41   template <typename T> struct WrappingIterator {
42     using value_type = T;
43 
44     const ArrayRef<value_type> Range;
45     typename decltype(Range)::const_iterator Position;
46 
47     // Rewind the tape, placing the position to again point at the beginning.
48     void rewind() { Position = Range.begin(); }
49 
50     // Advance position forward, possibly wrapping to the beginning.
51     // Returns whether the wrap happened.
52     bool advance() {
53       ++Position;
54       bool Wrapped = Position == Range.end();
55       if (Wrapped)
56         rewind();
57       return Wrapped;
58     }
59 
60     // Get the value at which we are currently pointing.
61     const value_type &operator*() const { return *Position; }
62 
63     WrappingIterator(ArrayRef<value_type> Range_) : Range(Range_) {
64       assert(!Range.empty() && "The range must not be empty.");
65       rewind();
66     }
67   };
68 
69   const ArrayRef<choices_storage_type> VariablesChoices;
70 
71   void performGeneration(
72       const function_ref<bool(ArrayRef<choice_type>)> Callback) const {
73     SmallVector<WrappingIterator<choice_type>, variable_smallsize>
74         VariablesState;
75 
76     // 'increment' of the the whole VariablesState is defined identically to the
77     // increment of a number: starting from the least significant element,
78     // increment it, and if it wrapped, then propagate that carry by also
79     // incrementing next (more significant) element.
80     auto IncrementState =
81         [](MutableArrayRef<WrappingIterator<choice_type>> VariablesState)
82         -> bool {
83       for (WrappingIterator<choice_type> &Variable :
84            llvm::reverse(VariablesState)) {
85         bool Wrapped = Variable.advance();
86         if (!Wrapped)
87           return false; // There you go, next combination is ready.
88         // We have carry - increment more significant variable next..
89       }
90       return true; // MSB variable wrapped, no more unique combinations.
91     };
92 
93     // Initialize the per-variable state to refer to the possible choices for
94     // that variable.
95     VariablesState.reserve(VariablesChoices.size());
96     for (ArrayRef<choice_type> VC : VariablesChoices)
97       VariablesState.emplace_back(VC);
98 
99     // Temporary buffer to store each combination before performing Callback.
100     SmallVector<choice_type, variable_smallsize> CurrentCombination;
101     CurrentCombination.resize(VariablesState.size());
102 
103     while (true) {
104       // Gather the currently-selected variable choices into a vector.
105       for (auto I : llvm::zip(VariablesState, CurrentCombination))
106         std::get<1>(I) = *std::get<0>(I);
107       // And pass the new combination into callback, as intended.
108       if (/*Abort=*/Callback(CurrentCombination))
109         return;
110       // And tick the state to next combination, which will be unique.
111       if (IncrementState(VariablesState))
112         return; // All combinations produced.
113     }
114   };
115 
116 public:
117   CombinationGenerator(ArrayRef<choices_storage_type> VariablesChoices_)
118       : VariablesChoices(VariablesChoices_) {
119 #ifndef NDEBUG
120     assert(!VariablesChoices.empty() && "There should be some variables.");
121     llvm::for_each(VariablesChoices, [](ArrayRef<choice_type> VariableChoices) {
122       assert(!VariableChoices.empty() &&
123              "There must always be some choice, at least a placeholder one.");
124     });
125 #endif
126   }
127 
128   // How many combinations can we produce, max?
129   // This is at most how many times the callback will be called.
130   size_t numCombinations() const {
131     size_t NumVariants = 1;
132     for (ArrayRef<choice_type> VariableChoices : VariablesChoices)
133       NumVariants *= VariableChoices.size();
134     assert(NumVariants >= 1 &&
135            "We should always end up producing at least one combination");
136     return NumVariants;
137   }
138 
139   // Actually perform exhaustive combination generation.
140   // Each result will be passed into the callback.
141   void generate(const function_ref<bool(ArrayRef<choice_type>)> Callback) {
142     performGeneration(Callback);
143   }
144 };
145 
146 } // namespace llvm
147 
148 #endif
149