Delete cost and loss functions when not in use.

Delete CostFunctions and LossFunctions when there are no more
ResidualBlocks referring to them. This is done by maintaining
a map with reference counts for CostFunctions and
LossFunctions.

The same maps are also used at the time of the destruction
of the ProblemImpl object itself. Previously vectors of these
objects were constructed, uniqed and the objects destroyed.

The update to the maps increases the cost of calling AddResidualBlock,
this has been mitigated, actually making AddResidualBlock faster, by
reusing a temporary vector rather than allocating one on the stack
every time.

Change-Id: I28b5287511713d28069ae428e2ff69224c0d03b4
diff --git a/internal/ceres/problem_impl.cc b/internal/ceres/problem_impl.cc
index 4088079..07da550 100644
--- a/internal/ceres/problem_impl.cc
+++ b/internal/ceres/problem_impl.cc
@@ -86,6 +86,27 @@
       << "size " << new_block_size << ".";
 }
 
+template <typename KeyType>
+void DecrementValueOrDeleteKey(const KeyType key,
+                               std::map<KeyType, int>* container) {
+  typename std::map<KeyType, int>::iterator it = container->find(key);
+  if (it->second == 1) {
+    delete key;
+    container->erase(it);
+  } else {
+    --it->second;
+  }
+}
+
+template <typename ForwardIterator>
+void STLDeleteContainerPairFirstPointers(ForwardIterator begin,
+                                         ForwardIterator end) {
+  while (begin != end) {
+    delete begin->first;
+    ++begin;
+  }
+}
+
 }  // namespace
 
 ParameterBlock* ProblemImpl::InternalAddParameterBlock(double* values,
@@ -177,16 +198,19 @@
   // The const casts here are legit, since ResidualBlock holds these
   // pointers as const pointers but we have ownership of them and
   // have the right to destroy them when the destructor is called.
-  if (options_.cost_function_ownership == TAKE_OWNERSHIP &&
-      residual_block->cost_function() != NULL) {
-    cost_functions_to_delete_.push_back(
-        const_cast<CostFunction*>(residual_block->cost_function()));
+  CostFunction* cost_function =
+      const_cast<CostFunction*>(residual_block->cost_function());
+  if (options_.cost_function_ownership == TAKE_OWNERSHIP) {
+    DecrementValueOrDeleteKey(cost_function, &cost_function_ref_count_);
   }
+
+  LossFunction* loss_function =
+      const_cast<LossFunction*>(residual_block->loss_function());
   if (options_.loss_function_ownership == TAKE_OWNERSHIP &&
-      residual_block->loss_function() != NULL) {
-    loss_functions_to_delete_.push_back(
-        const_cast<LossFunction*>(residual_block->loss_function()));
+      loss_function != NULL) {
+    DecrementValueOrDeleteKey(loss_function, &loss_function_ref_count_);
   }
+
   delete residual_block;
 }
 
@@ -205,18 +229,28 @@
   delete parameter_block;
 }
 
-ProblemImpl::ProblemImpl() : program_(new internal::Program) {}
+ProblemImpl::ProblemImpl()
+    : program_(new internal::Program) {
+  residual_parameters_.reserve(10);
+}
+
 ProblemImpl::ProblemImpl(const Problem::Options& options)
-    : options_(options),
-      program_(new internal::Program) {}
+    : options_(options), program_(new internal::Program) {
+  residual_parameters_.reserve(10);
+}
 
 ProblemImpl::~ProblemImpl() {
-  // Collect the unique cost/loss functions and delete the residuals.
-  const int num_residual_blocks = program_->residual_blocks_.size();
-  cost_functions_to_delete_.reserve(num_residual_blocks);
-  loss_functions_to_delete_.reserve(num_residual_blocks);
-  for (int i = 0; i < program_->residual_blocks_.size(); ++i) {
-    DeleteBlock(program_->residual_blocks_[i]);
+  STLDeleteContainerPointers(program_->residual_blocks_.begin(),
+                             program_->residual_blocks_.end());
+
+  if (options_.cost_function_ownership == TAKE_OWNERSHIP) {
+    STLDeleteContainerPairFirstPointers(cost_function_ref_count_.begin(),
+                                        cost_function_ref_count_.end());
+  }
+
+  if (options_.loss_function_ownership == TAKE_OWNERSHIP) {
+    STLDeleteContainerPairFirstPointers(loss_function_ref_count_.begin(),
+                                        loss_function_ref_count_.end());
   }
 
   // Collect the unique parameterizations and delete the parameters.
@@ -224,13 +258,9 @@
     DeleteBlock(program_->parameter_blocks_[i]);
   }
 
-  // Delete the owned cost/loss functions and parameterizations.
+  // Delete the owned parameterizations.
   STLDeleteUniqueContainerPointers(local_parameterizations_to_delete_.begin(),
                                    local_parameterizations_to_delete_.end());
-  STLDeleteUniqueContainerPointers(cost_functions_to_delete_.begin(),
-                                   cost_functions_to_delete_.end());
-  STLDeleteUniqueContainerPointers(loss_functions_to_delete_.begin(),
-                                   loss_functions_to_delete_.end());
 }
 
 ResidualBlock* ProblemImpl::AddResidualBlock(
@@ -309,6 +339,15 @@
     residual_block_set_.insert(new_residual_block);
   }
 
+  if (options_.cost_function_ownership == TAKE_OWNERSHIP) {
+    ++cost_function_ref_count_[cost_function];
+  }
+
+  if (options_.loss_function_ownership == TAKE_OWNERSHIP &&
+      loss_function != NULL) {
+    ++loss_function_ref_count_[loss_function];
+  }
+
   return new_residual_block;
 }
 
@@ -318,69 +357,69 @@
     CostFunction* cost_function,
     LossFunction* loss_function,
     double* x0) {
-  vector<double*> residual_parameters;
-  residual_parameters.push_back(x0);
-  return AddResidualBlock(cost_function, loss_function, residual_parameters);
+  residual_parameters_.clear();
+  residual_parameters_.push_back(x0);
+  return AddResidualBlock(cost_function, loss_function, residual_parameters_);
 }
 
 ResidualBlock* ProblemImpl::AddResidualBlock(
     CostFunction* cost_function,
     LossFunction* loss_function,
     double* x0, double* x1) {
-  vector<double*> residual_parameters;
-  residual_parameters.push_back(x0);
-  residual_parameters.push_back(x1);
-  return AddResidualBlock(cost_function, loss_function, residual_parameters);
+  residual_parameters_.clear();
+  residual_parameters_.push_back(x0);
+  residual_parameters_.push_back(x1);
+  return AddResidualBlock(cost_function, loss_function, residual_parameters_);
 }
 
 ResidualBlock* ProblemImpl::AddResidualBlock(
     CostFunction* cost_function,
     LossFunction* loss_function,
     double* x0, double* x1, double* x2) {
-  vector<double*> residual_parameters;
-  residual_parameters.push_back(x0);
-  residual_parameters.push_back(x1);
-  residual_parameters.push_back(x2);
-  return AddResidualBlock(cost_function, loss_function, residual_parameters);
+  residual_parameters_.clear();
+  residual_parameters_.push_back(x0);
+  residual_parameters_.push_back(x1);
+  residual_parameters_.push_back(x2);
+  return AddResidualBlock(cost_function, loss_function, residual_parameters_);
 }
 
 ResidualBlock* ProblemImpl::AddResidualBlock(
     CostFunction* cost_function,
     LossFunction* loss_function,
     double* x0, double* x1, double* x2, double* x3) {
-  vector<double*> residual_parameters;
-  residual_parameters.push_back(x0);
-  residual_parameters.push_back(x1);
-  residual_parameters.push_back(x2);
-  residual_parameters.push_back(x3);
-  return AddResidualBlock(cost_function, loss_function, residual_parameters);
+  residual_parameters_.clear();
+  residual_parameters_.push_back(x0);
+  residual_parameters_.push_back(x1);
+  residual_parameters_.push_back(x2);
+  residual_parameters_.push_back(x3);
+  return AddResidualBlock(cost_function, loss_function, residual_parameters_);
 }
 
 ResidualBlock* ProblemImpl::AddResidualBlock(
     CostFunction* cost_function,
     LossFunction* loss_function,
     double* x0, double* x1, double* x2, double* x3, double* x4) {
-  vector<double*> residual_parameters;
-  residual_parameters.push_back(x0);
-  residual_parameters.push_back(x1);
-  residual_parameters.push_back(x2);
-  residual_parameters.push_back(x3);
-  residual_parameters.push_back(x4);
-  return AddResidualBlock(cost_function, loss_function, residual_parameters);
+  residual_parameters_.clear();
+  residual_parameters_.push_back(x0);
+  residual_parameters_.push_back(x1);
+  residual_parameters_.push_back(x2);
+  residual_parameters_.push_back(x3);
+  residual_parameters_.push_back(x4);
+  return AddResidualBlock(cost_function, loss_function, residual_parameters_);
 }
 
 ResidualBlock* ProblemImpl::AddResidualBlock(
     CostFunction* cost_function,
     LossFunction* loss_function,
     double* x0, double* x1, double* x2, double* x3, double* x4, double* x5) {
-  vector<double*> residual_parameters;
-  residual_parameters.push_back(x0);
-  residual_parameters.push_back(x1);
-  residual_parameters.push_back(x2);
-  residual_parameters.push_back(x3);
-  residual_parameters.push_back(x4);
-  residual_parameters.push_back(x5);
-  return AddResidualBlock(cost_function, loss_function, residual_parameters);
+  residual_parameters_.clear();
+  residual_parameters_.push_back(x0);
+  residual_parameters_.push_back(x1);
+  residual_parameters_.push_back(x2);
+  residual_parameters_.push_back(x3);
+  residual_parameters_.push_back(x4);
+  residual_parameters_.push_back(x5);
+  return AddResidualBlock(cost_function, loss_function, residual_parameters_);
 }
 
 ResidualBlock* ProblemImpl::AddResidualBlock(
@@ -388,15 +427,15 @@
     LossFunction* loss_function,
     double* x0, double* x1, double* x2, double* x3, double* x4, double* x5,
     double* x6) {
-  vector<double*> residual_parameters;
-  residual_parameters.push_back(x0);
-  residual_parameters.push_back(x1);
-  residual_parameters.push_back(x2);
-  residual_parameters.push_back(x3);
-  residual_parameters.push_back(x4);
-  residual_parameters.push_back(x5);
-  residual_parameters.push_back(x6);
-  return AddResidualBlock(cost_function, loss_function, residual_parameters);
+  residual_parameters_.clear();
+  residual_parameters_.push_back(x0);
+  residual_parameters_.push_back(x1);
+  residual_parameters_.push_back(x2);
+  residual_parameters_.push_back(x3);
+  residual_parameters_.push_back(x4);
+  residual_parameters_.push_back(x5);
+  residual_parameters_.push_back(x6);
+  return AddResidualBlock(cost_function, loss_function, residual_parameters_);
 }
 
 ResidualBlock* ProblemImpl::AddResidualBlock(
@@ -404,16 +443,16 @@
     LossFunction* loss_function,
     double* x0, double* x1, double* x2, double* x3, double* x4, double* x5,
     double* x6, double* x7) {
-  vector<double*> residual_parameters;
-  residual_parameters.push_back(x0);
-  residual_parameters.push_back(x1);
-  residual_parameters.push_back(x2);
-  residual_parameters.push_back(x3);
-  residual_parameters.push_back(x4);
-  residual_parameters.push_back(x5);
-  residual_parameters.push_back(x6);
-  residual_parameters.push_back(x7);
-  return AddResidualBlock(cost_function, loss_function, residual_parameters);
+  residual_parameters_.clear();
+  residual_parameters_.push_back(x0);
+  residual_parameters_.push_back(x1);
+  residual_parameters_.push_back(x2);
+  residual_parameters_.push_back(x3);
+  residual_parameters_.push_back(x4);
+  residual_parameters_.push_back(x5);
+  residual_parameters_.push_back(x6);
+  residual_parameters_.push_back(x7);
+  return AddResidualBlock(cost_function, loss_function, residual_parameters_);
 }
 
 ResidualBlock* ProblemImpl::AddResidualBlock(
@@ -421,17 +460,17 @@
     LossFunction* loss_function,
     double* x0, double* x1, double* x2, double* x3, double* x4, double* x5,
     double* x6, double* x7, double* x8) {
-  vector<double*> residual_parameters;
-  residual_parameters.push_back(x0);
-  residual_parameters.push_back(x1);
-  residual_parameters.push_back(x2);
-  residual_parameters.push_back(x3);
-  residual_parameters.push_back(x4);
-  residual_parameters.push_back(x5);
-  residual_parameters.push_back(x6);
-  residual_parameters.push_back(x7);
-  residual_parameters.push_back(x8);
-  return AddResidualBlock(cost_function, loss_function, residual_parameters);
+  residual_parameters_.clear();
+  residual_parameters_.push_back(x0);
+  residual_parameters_.push_back(x1);
+  residual_parameters_.push_back(x2);
+  residual_parameters_.push_back(x3);
+  residual_parameters_.push_back(x4);
+  residual_parameters_.push_back(x5);
+  residual_parameters_.push_back(x6);
+  residual_parameters_.push_back(x7);
+  residual_parameters_.push_back(x8);
+  return AddResidualBlock(cost_function, loss_function, residual_parameters_);
 }
 
 ResidualBlock* ProblemImpl::AddResidualBlock(
@@ -439,18 +478,18 @@
     LossFunction* loss_function,
     double* x0, double* x1, double* x2, double* x3, double* x4, double* x5,
     double* x6, double* x7, double* x8, double* x9) {
-  vector<double*> residual_parameters;
-  residual_parameters.push_back(x0);
-  residual_parameters.push_back(x1);
-  residual_parameters.push_back(x2);
-  residual_parameters.push_back(x3);
-  residual_parameters.push_back(x4);
-  residual_parameters.push_back(x5);
-  residual_parameters.push_back(x6);
-  residual_parameters.push_back(x7);
-  residual_parameters.push_back(x8);
-  residual_parameters.push_back(x9);
-  return AddResidualBlock(cost_function, loss_function, residual_parameters);
+  residual_parameters_.clear();
+  residual_parameters_.push_back(x0);
+  residual_parameters_.push_back(x1);
+  residual_parameters_.push_back(x2);
+  residual_parameters_.push_back(x3);
+  residual_parameters_.push_back(x4);
+  residual_parameters_.push_back(x5);
+  residual_parameters_.push_back(x6);
+  residual_parameters_.push_back(x7);
+  residual_parameters_.push_back(x8);
+  residual_parameters_.push_back(x9);
+  return AddResidualBlock(cost_function, loss_function, residual_parameters_);
 }
 
 void ProblemImpl::AddParameterBlock(double* values, int size) {
diff --git a/internal/ceres/problem_impl.h b/internal/ceres/problem_impl.h
index a4689c3..b7f258c 100644
--- a/internal/ceres/problem_impl.h
+++ b/internal/ceres/problem_impl.h
@@ -65,6 +65,8 @@
  public:
   typedef std::map<double*, ParameterBlock*> ParameterMap;
   typedef HashSet<ResidualBlock*> ResidualBlockSet;
+  typedef std::map<CostFunction*, int> CostFunctionRefCount;
+  typedef std::map<LossFunction*, int> LossFunctionRefCount;
 
   ProblemImpl();
   explicit ProblemImpl(const Problem::Options& options);
@@ -209,16 +211,21 @@
   // The actual parameter and residual blocks.
   internal::scoped_ptr<internal::Program> program_;
 
-  // When removing residual and parameter blocks, cost/loss functions and
-  // parameterizations have ambiguous ownership. Instead of scanning the entire
-  // problem to see if the cost/loss/parameterization is shared with other
-  // residual or parameter blocks, buffer them until destruction.
+  // When removing parameter blocks, parameterizations have ambiguous
+  // ownership. Instead of scanning the entire problem to see if the
+  // parameterization is shared with other parameter blocks, buffer
+  // them until destruction.
   //
   // TODO(keir): See if it makes sense to use sets instead.
-  std::vector<CostFunction*> cost_functions_to_delete_;
-  std::vector<LossFunction*> loss_functions_to_delete_;
   std::vector<LocalParameterization*> local_parameterizations_to_delete_;
 
+  // For each cost function and loss function in the problem, a count
+  // of the number of residual blocks that refer to them. When the
+  // count goes to zero and the problem owns these objects, they are
+  // destroyed.
+  CostFunctionRefCount cost_function_ref_count_;
+  LossFunctionRefCount loss_function_ref_count_;
+  std::vector<double*> residual_parameters_;
   CERES_DISALLOW_COPY_AND_ASSIGN(ProblemImpl);
 };
 
diff --git a/internal/ceres/problem_test.cc b/internal/ceres/problem_test.cc
index 08d4b04..826e2c2 100644
--- a/internal/ceres/problem_test.cc
+++ b/internal/ceres/problem_test.cc
@@ -385,11 +385,10 @@
     ResidualBlock* r_wz = problem.AddResidualBlock(cost_wz, NULL, w, z);
     EXPECT_EQ(2, problem.NumResidualBlocks());
 
-    // In the current implementation, the destructor shouldn't get run yet.
     problem.RemoveResidualBlock(r_yz);
-    CHECK_EQ(num_destructions, 0);
+    CHECK_EQ(num_destructions, 1);
     problem.RemoveResidualBlock(r_wz);
-    CHECK_EQ(num_destructions, 0);
+    CHECK_EQ(num_destructions, 2);
 
     EXPECT_EQ(0, problem.NumResidualBlocks());
   }
diff --git a/internal/ceres/residual_block.cc b/internal/ceres/residual_block.cc
index 9a123cf..f73d54f 100644
--- a/internal/ceres/residual_block.cc
+++ b/internal/ceres/residual_block.cc
@@ -54,7 +54,7 @@
     const LossFunction* loss_function,
     const std::vector<ParameterBlock*>& parameter_blocks,
     int index)
-    : cost_function_(cost_function),
+    : cost_function_(CHECK_NOTNULL(cost_function)),
       loss_function_(loss_function),
       parameter_blocks_(
           new ParameterBlock* [