Bookeep ExpressionPtr with CapacityConstraint.

schedulers/tetrisched/include/tetrisched/CapacityConstraint.hpp

@@ -20,6 +20,9 @@ class CapacityConstraint {
   /// The SolverModel constraint that enforces the resource usage.
   ConstraintPtr capacityConstraint;
 
+  /// A Map from the ExpressionPtr to the usage of that Expression.
+  std::unordered_map<ExpressionPtr, XOrVariableT<uint32_t>> usageMap;
+
   /// The CapacityConstraintMap is allowed to translate this CapacityConstraint.
   void translate(SolverModelPtr solverModel);
   friend class CapacityConstraintMap;
@@ -30,8 +33,8 @@ class CapacityConstraint {
   CapacityConstraint(const Partition& partition, Time time);
 
   /// Registers the given usage in this CapacityConstraint.
-  void registerUsage(uint32_t usage);
-  void registerUsage(VariablePtr variable);
+  void registerUsage(const ExpressionPtr expression, uint32_t usage);
+  void registerUsage(const ExpressionPtr expression, VariablePtr variable);
 };
 using CapacityConstraintPtr = std::shared_ptr<CapacityConstraint>;
 
@@ -60,32 +63,36 @@ class CapacityConstraintMap {
   /// Initialize a CapacityConstraintMap with the granularity of 1.
   CapacityConstraintMap();
 
-  /// Registers the usage for the given Partition at the given time
-  /// as specified by the value of the variable, which is to be
+  /// Registers the usage by the Expression for the Partition at the
+  /// time specified by the value of the variable, which is to be
   /// decided by the solver.
-  void registerUsageAtTime(const Partition& partition, Time time,
+  void registerUsageAtTime(const ExpressionPtr expression,
+                           const Partition& partition, Time time,
                            VariablePtr variable);
 
-  /// Registers the usage for the given Partition at the given time
-  /// as specified by the known usage.
-  void registerUsageAtTime(const Partition& partition, Time time,
+  /// Registers the usage by the Expression for the Partition at the
+  /// time specified by the known usage.
+  void registerUsageAtTime(const ExpressionPtr expression,
+                           const Partition& partition, Time time,
                            uint32_t usage);
 
-  /// Registers the usage for the given Partition in the time range
-  /// starting from startTime and lasting for duration as specified
-  /// by the value of the variable, which is to be decided by the solver.
+  /// Registers the usage by the Expression for the Partition in the
+  /// time range starting from startTime and lasting for duration as
+  /// specified by the variable, which is to be decided by the solver.
   /// Optionally, a step granularity can be provided. The default granularity
   /// is the one that the CapacityConstraintMap was initialized with.
-  void registerUsageForDuration(const Partition& partition, Time startTime,
+  void registerUsageForDuration(const ExpressionPtr expression,
+                                const Partition& partition, Time startTime,
                                 Time duration, VariablePtr variable,
                                 std::optional<Time> granularity);
 
-  /// Registers the usage for the given Partition in the time range
-  /// starting from startTime and lasting for duration as specified
-  /// by the value of the variable known at runtime. Optionally, a step
+  /// Registers the usage by the Expression for the Partition in the
+  /// time range starting from startTime and lasting for duration as
+  /// specified by the variable known at runtime. Optionally, a step
   /// granularity can be provided. The default granularity is the one
   /// that the CapacityConstraintMap was initialized with.
-  void registerUsageForDuration(const Partition& partition, Time startTime,
+  void registerUsageForDuration(const ExpressionPtr expression,
+                                const Partition& partition, Time startTime,
                                 Time duration, uint32_t usage,
                                 std::optional<Time> granularity);
 

schedulers/tetrisched/src/CapacityConstraint.cpp

@@ -21,16 +21,20 @@ CapacityConstraint::CapacityConstraint(const Partition& partition, Time time)
               std::to_string(time),
           ConstraintType::CONSTR_LE, partition.getQuantity())) {}
 
-void CapacityConstraint::registerUsage(uint32_t usage) {
+void CapacityConstraint::registerUsage(const ExpressionPtr expression,
+                                       uint32_t usage) {
   if (usage == 0) {
     // No usage was registered. We don't need to add anything.
     return;
   }
   capacityConstraint->addTerm(usage);
+  usageMap[expression] = usage;
 }
 
-void CapacityConstraint::registerUsage(VariablePtr variable) {
+void CapacityConstraint::registerUsage(const ExpressionPtr expression,
+                                       VariablePtr variable) {
   capacityConstraint->addTerm(variable);
+  usageMap[expression] = variable;
 }
 
 void CapacityConstraint::translate(SolverModelPtr solverModel) {
@@ -52,7 +56,8 @@ CapacityConstraintMap::CapacityConstraintMap(Time granularity)
 
 CapacityConstraintMap::CapacityConstraintMap() : granularity(1) {}
 
-void CapacityConstraintMap::registerUsageAtTime(const Partition& partition,
+void CapacityConstraintMap::registerUsageAtTime(const ExpressionPtr expression,
+                                                const Partition& partition,
                                                 Time time,
                                                 VariablePtr variable) {
   // Get or insert the Constraint corresponding to this partition and time.
@@ -63,10 +68,11 @@ void CapacityConstraintMap::registerUsageAtTime(const Partition& partition,
   }
 
   // Add the variable to the Constraint.
-  capacityConstraints[mapKey]->registerUsage(variable);
+  capacityConstraints[mapKey]->registerUsage(expression, variable);
 }
 
-void CapacityConstraintMap::registerUsageAtTime(const Partition& partition,
+void CapacityConstraintMap::registerUsageAtTime(const ExpressionPtr expression,
+                                                const Partition& partition,
                                                 Time time, uint32_t usage) {
   if (usage == 0) {
     // No usage was registered. We don't need to add anything.
@@ -80,26 +86,26 @@ void CapacityConstraintMap::registerUsageAtTime(const Partition& partition,
   }
 
   // Add the variable to the Constraint.
-  capacityConstraints[mapKey]->registerUsage(usage);
+  capacityConstraints[mapKey]->registerUsage(expression, usage);
 }
 
 void CapacityConstraintMap::registerUsageForDuration(
-    const Partition& partition, Time startTime, Time duration,
-    VariablePtr variable, std::optional<Time> granularity) {
+    const ExpressionPtr expression, const Partition& partition, Time startTime,
+    Time duration, VariablePtr variable, std::optional<Time> granularity) {
   Time _granularity = granularity.value_or(this->granularity);
   for (Time time = startTime; time < startTime + duration;
        time += _granularity) {
-    registerUsageAtTime(partition, time, variable);
+    registerUsageAtTime(expression, partition, time, variable);
   }
 }
 
 void CapacityConstraintMap::registerUsageForDuration(
-    const Partition& partition, Time startTime, Time duration, uint32_t usage,
-    std::optional<Time> granularity) {
+    const ExpressionPtr expression, const Partition& partition, Time startTime,
+    Time duration, uint32_t usage, std::optional<Time> granularity) {
   Time _granularity = granularity.value_or(this->granularity);
   for (Time time = startTime; time < startTime + duration;
        time += _granularity) {
-    registerUsageAtTime(partition, time, usage);
+    registerUsageAtTime(expression, partition, time, usage);
   }
 }
 

schedulers/tetrisched/src/Expression.cpp

@@ -441,8 +441,9 @@ ParseResultPtr ChooseExpression::parse(
 
     // Register this indicator with the capacity constraints that
     // are being bubbled up.
-    capacityConstraints.registerUsageForDuration(
-        *partition, startTime, duration, allocationVar, std::nullopt);
+    capacityConstraints.registerUsageForDuration(shared_from_this(), *partition,
+                                                 startTime, duration,
+                                                 allocationVar, std::nullopt);
   }
   // Ensure that if the Choose expression is satisfied, it fulfills the
   // demand for this expression. Pass the constraint to the model.
@@ -594,7 +595,8 @@ ParseResultPtr MalleableChooseExpression::parse(
         // Register this Integer variable with the CapacityConstraintMap
         // that is being bubbled up.
         capacityConstraints.registerUsageForDuration(
-            *partition, time, granularity, allocationAtTime, std::nullopt);
+            shared_from_this(), *partition, time, granularity, allocationAtTime,
+            std::nullopt);
       } else {
         throw tetrisched::exceptions::ExpressionConstructionException(
             "Multiple variables detected for the Partition " +
@@ -895,9 +897,9 @@ ParseResultPtr AllocationExpression::parse(
       std::make_shared<ObjectiveFunction>(ObjectiveType::OBJ_MAXIMIZE);
   (parsedResult->utility).value()->addTerm(1);
   for (const auto& [partition, allocation] : allocatedResources) {
-    capacityConstraints.registerUsageForDuration(
-        *partition, startTime, duration, static_cast<uint32_t>(allocation),
-        std::nullopt);
+    capacityConstraints.registerUsageForDuration(shared_from_this(), *partition,
+                                                 startTime, duration,
+                                                 allocation, std::nullopt);
   }
   return parsedResult;
 }