PRFSet implementation: necessary subtle changes to support PRFSet.

PiperOrigin-RevId: 303407705

2 files changed, 9 insertions, 7 deletions

diff --git a/cc/prf/prf_set.h b/cc/prf/prf_set.h
index 17db0ff5e..70055c950 100644
--- a/cc/prf/prf_set.h
+++ b/cc/prf/prf_set.h
@@ -69,7 +69,7 @@ class PrfSet {
   virtual uint32_t GetPrimaryId() const = 0;
   // A map of the PRFs represented by the keys in this keyset.
   // The map is guaranteed to contain getPrimaryId() as a key.
-  virtual std::map<uint32_t, Prf*> GetPrfs() const = 0;
+  virtual const std::map<uint32_t, Prf*>& GetPrfs() const = 0;
   // Convenience method to compute the primary PRF on a given input.
   // See PRF.compute for details of the parameters.
   util::StatusOr<std::string> ComputePrimary(absl::string_view input,
diff --git a/cc/prf/prf_set_test.cc b/cc/prf/prf_set_test.cc
index b80cc7477..e7ee43981 100644
--- a/cc/prf/prf_set_test.cc
+++ b/cc/prf/prf_set_test.cc
@@ -36,10 +36,10 @@ class DummyPrf : public Prf {
 class DummyPrfSet : public PrfSet {
  public:
   uint32_t GetPrimaryId() const override { return 1; }
-  std::map<uint32_t, Prf*> GetPrfs() const override {
-    std::map<uint32_t, Prf*> prfs;
-    prfs.insert({1, dummy_.get()});
-    return prfs;
+  const std::map<uint32_t, Prf*>& GetPrfs() const override {
+    static const std::map<uint32_t, Prf*>* prfs =
+        new std::map<uint32_t, Prf*>({{1, dummy_.get()}});
+    return *prfs;
   }
 
  private:
@@ -49,8 +49,10 @@ class DummyPrfSet : public PrfSet {
 class BrokenDummyPrfSet : public PrfSet {
  public:
   uint32_t GetPrimaryId() const override { return 1; }
-  std::map<uint32_t, Prf*> GetPrfs() const override {
-    return std::map<uint32_t, Prf*>();
+  const std::map<uint32_t, Prf*>& GetPrfs() const override {
+    static const std::map<uint32_t, Prf*>* prfs =
+        new std::map<uint32_t, Prf*>();
+    return *prfs;
   }
 };