1 files changed, 58 insertions, 7 deletions

diff --git a/internal/ceres/parameter_block.h b/internal/ceres/parameter_block.h
index f20805c..695fa6f 100644
--- a/internal/ceres/parameter_block.h
+++ b/internal/ceres/parameter_block.h
@@ -1,5 +1,5 @@
 // Ceres Solver - A fast non-linear least squares minimizer
-// Copyright 2010, 2011, 2012 Google Inc. All rights reserved.
+// Copyright 2010, 2011, 2012, 2013 Google Inc. All rights reserved.
 // http://code.google.com/p/ceres-solver/
 //
 // Redistribution and use in source and binary forms, with or without
@@ -34,6 +34,7 @@
 #include <cstdlib>
 #include <string>
 #include "ceres/array_utils.h"
+#include "ceres/collections_port.h"
 #include "ceres/integral_types.h"
 #include "ceres/internal/eigen.h"
 #include "ceres/internal/port.h"
@@ -46,6 +47,7 @@ namespace ceres {
 namespace internal {
 
 class ProblemImpl;
+class ResidualBlock;
 
 // The parameter block encodes the location of the user's original value, and
 // also the "current state" of the parameter. The evaluator uses whatever is in
@@ -58,13 +60,29 @@ class ProblemImpl;
 // responsible for the proper disposal of the local parameterization.
 class ParameterBlock {
  public:
-  ParameterBlock(double* user_state, int size) {
-    Init(user_state, size, NULL);
+  // TODO(keir): Decide what data structure is best here. Should this be a set?
+  // Probably not, because sets are memory inefficient. However, if it's a
+  // vector, you can get into pathological linear performance when removing a
+  // residual block from a problem where all the residual blocks depend on one
+  // parameter; for example, shared focal length in a bundle adjustment
+  // problem. It might be worth making a custom structure that is just an array
+  // when it is small, but transitions to a hash set when it has more elements.
+  //
+  // For now, use a hash set.
+  typedef HashSet<ResidualBlock*> ResidualBlockSet;
+
+  // Create a parameter block with the user state, size, and index specified.
+  // The size is the size of the parameter block and the index is the position
+  // of the parameter block inside a Program (if any).
+  ParameterBlock(double* user_state, int size, int index) {
+    Init(user_state, size, index, NULL);
   }
+
   ParameterBlock(double* user_state,
                  int size,
+                 int index,
                  LocalParameterization* local_parameterization) {
-    Init(user_state, size, local_parameterization);
+    Init(user_state, size, index, local_parameterization);
   }
 
   // The size of the parameter block.
@@ -155,8 +173,8 @@ class ParameterBlock {
           new double[local_parameterization_->GlobalSize() *
                      local_parameterization_->LocalSize()]);
       CHECK(UpdateLocalParameterizationJacobian())
-          "Local parameterization Jacobian computation failed"
-          "for x: " << ConstVectorRef(state_, Size()).transpose();
+          << "Local parameterization Jacobian computation failed for x: "
+          << ConstVectorRef(state_, Size()).transpose();
     } else {
       // Ignore the case that the parameterizations match.
     }
@@ -187,12 +205,43 @@ class ParameterBlock {
                         delta_offset_);
   }
 
+  void EnableResidualBlockDependencies() {
+    CHECK(residual_blocks_.get() == NULL)
+        << "Ceres bug: There is already a residual block collection "
+        << "for parameter block: " << ToString();
+    residual_blocks_.reset(new ResidualBlockSet);
+  }
+
+  void AddResidualBlock(ResidualBlock* residual_block) {
+    CHECK(residual_blocks_.get() != NULL)
+        << "Ceres bug: The residual block collection is null for parameter "
+        << "block: " << ToString();
+    residual_blocks_->insert(residual_block);
+  }
+
+  void RemoveResidualBlock(ResidualBlock* residual_block) {
+    CHECK(residual_blocks_.get() != NULL)
+        << "Ceres bug: The residual block collection is null for parameter "
+        << "block: " << ToString();
+    CHECK(residual_blocks_->find(residual_block) != residual_blocks_->end())
+        << "Ceres bug: Missing residual for parameter block: " << ToString();
+    residual_blocks_->erase(residual_block);
+  }
+
+  // This is only intended for iterating; perhaps this should only expose
+  // .begin() and .end().
+  ResidualBlockSet* mutable_residual_blocks() {
+    return residual_blocks_.get();
+  }
+
  private:
   void Init(double* user_state,
             int size,
+            int index,
             LocalParameterization* local_parameterization) {
     user_state_ = user_state;
     size_ = size;
+    index_ = index;
     is_constant_ = false;
     state_ = user_state_;
 
@@ -201,7 +250,6 @@ class ParameterBlock {
       SetParameterization(local_parameterization);
     }
 
-    index_ = -1;
     state_offset_ = -1;
     delta_offset_ = -1;
   }
@@ -261,6 +309,9 @@ class ParameterBlock {
   // The offset of this parameter block inside a larger delta vector.
   int32 delta_offset_;
 
+  // If non-null, contains the residual blocks this parameter block is in.
+  scoped_ptr<ResidualBlockSet> residual_blocks_;
+
   // Necessary so ProblemImpl can clean up the parameterizations.
   friend class ProblemImpl;
 };