Various changes to JBig2 cache:

- Makes the cache be per-document
- Keys the cache on ObjNum and stream offset instead of keying on a pointer to the data (which can result in false cache hits).
- Makes it so the cache is only used for the globals stream.
- Reenable the cache.

R=thestig@chromium.org

BUG=pdfium:207

Review URL: https://codereview.chromium.org/1380243004 .

9 files changed, 119 insertions, 86 deletions

diff --git a/core/src/fpdfapi/fpdf_render/fpdf_render_loadimage.cpp b/core/src/fpdfapi/fpdf_render/fpdf_render_loadimage.cpp
index f1fbf41d6d..1e7f077d48 100644
--- a/core/src/fpdfapi/fpdf_render/fpdf_render_loadimage.cpp
+++ b/core/src/fpdfapi/fpdf_render/fpdf_render_loadimage.cpp
@@ -407,12 +407,10 @@ int CPDF_DIBSource::ContinueLoadDIBSource(IFX_Pause* pPause) {
           m_pGlobalStream->LoadAllData(pGlobals, FALSE);
         }
       }
-      ret = pJbig2Module->StartDecode(
-          m_pJbig2Context, m_Width, m_Height, m_pStreamAcc->GetData(),
-          m_pStreamAcc->GetSize(),
-          m_pGlobalStream ? m_pGlobalStream->GetData() : NULL,
-          m_pGlobalStream ? m_pGlobalStream->GetSize() : 0,
-          m_pCachedBitmap->GetBuffer(), m_pCachedBitmap->GetPitch(), pPause);
+      ret = pJbig2Module->StartDecode(m_pJbig2Context, m_pDocument, m_Width,
+                                      m_Height, m_pStreamAcc, m_pGlobalStream,
+                                      m_pCachedBitmap->GetBuffer(),
+                                      m_pCachedBitmap->GetPitch(), pPause);
       if (ret < 0) {
         m_pCachedBitmap.reset();
         delete m_pGlobalStream;
diff --git a/core/src/fxcodec/codec/codec_int.h b/core/src/fxcodec/codec/codec_int.h
index 10d0c3fd79..b08dee7422 100644
--- a/core/src/fxcodec/codec/codec_int.h
+++ b/core/src/fxcodec/codec/codec_int.h
@@ -274,10 +274,8 @@ class CCodec_Jbig2Context {
 
   FX_DWORD m_width;
   FX_DWORD m_height;
-  uint8_t* m_src_buf;
-  FX_DWORD m_src_size;
-  const uint8_t* m_global_data;
-  FX_DWORD m_global_size;
+  CPDF_StreamAcc* m_pGlobalStream;
+  CPDF_StreamAcc* m_pSrcStream;
   uint8_t* m_dest_buf;
   FX_DWORD m_dest_pitch;
   IFX_Pause* m_pPause;
@@ -292,21 +290,17 @@ class CCodec_Jbig2Module : public ICodec_Jbig2Module {
   // ICodec_Jbig2Module
   void* CreateJbig2Context() override;
   FXCODEC_STATUS StartDecode(void* pJbig2Context,
+                             CFX_PrivateData* pPrivateData,
                              FX_DWORD width,
                              FX_DWORD height,
-                             const uint8_t* src_buf,
-                             FX_DWORD src_size,
-                             const uint8_t* global_data,
-                             FX_DWORD global_size,
+                             CPDF_StreamAcc* src_stream,
+                             CPDF_StreamAcc* global_stream,
                              uint8_t* dest_buf,
                              FX_DWORD dest_pitch,
                              IFX_Pause* pPause) override;
   FXCODEC_STATUS ContinueDecode(void* pJbig2Context,
                                 IFX_Pause* pPause) override;
   void DestroyJbig2Context(void* pJbig2Context) override;
-
- private:
-  std::list<CJBig2_CachePair> m_SymbolDictCache;
 };
 
 struct DecodeData {
diff --git a/core/src/fxcodec/codec/fx_codec_jbig.cpp b/core/src/fxcodec/codec/fx_codec_jbig.cpp
index 18a665784a..1524b68df0 100644
--- a/core/src/fxcodec/codec/fx_codec_jbig.cpp
+++ b/core/src/fxcodec/codec/fx_codec_jbig.cpp
@@ -7,14 +7,41 @@
 #include "../../../include/fxcodec/fx_codec.h"
 #include "codec_int.h"
 
+// Holds per-document JBig2 related data.
+class JBig2DocumentContext : public CFX_DestructObject {
+ public:
+  std::list<CJBig2_CachePair>* GetSymbolDictCache() {
+    return &m_SymbolDictCache;
+  }
+
+  ~JBig2DocumentContext() {
+    for (auto it : m_SymbolDictCache) {
+      delete it.second;
+    }
+  }
+
+ private:
+  std::list<CJBig2_CachePair> m_SymbolDictCache;
+};
+
+JBig2DocumentContext* GetJBig2DocumentContext(CCodec_Jbig2Module* pModule,
+                                              CFX_PrivateData* pPrivateData) {
+  void* pModulePrivateData = pPrivateData->GetPrivateData(pModule);
+  if (pModulePrivateData) {
+    CFX_DestructObject* pDestructObject =
+        reinterpret_cast<CFX_DestructObject*>(pModulePrivateData);
+    return static_cast<JBig2DocumentContext*>(pDestructObject);
+  }
+  JBig2DocumentContext* pJBig2DocumentContext = new JBig2DocumentContext();
+  pPrivateData->SetPrivateObj(pModule, pJBig2DocumentContext);
+  return pJBig2DocumentContext;
+}
+
 CCodec_Jbig2Context::CCodec_Jbig2Context() {
   FXSYS_memset(this, 0, sizeof(CCodec_Jbig2Context));
 }
 
 CCodec_Jbig2Module::~CCodec_Jbig2Module() {
-  for (auto it : m_SymbolDictCache) {
-    delete it.second;
-  }
 }
 
 void* CCodec_Jbig2Module::CreateJbig2Context() {
@@ -29,31 +56,31 @@ void CCodec_Jbig2Module::DestroyJbig2Context(void* pJbig2Content) {
   pJbig2Content = NULL;
 }
 FXCODEC_STATUS CCodec_Jbig2Module::StartDecode(void* pJbig2Context,
+                                               CFX_PrivateData* pPrivateData,
                                                FX_DWORD width,
                                                FX_DWORD height,
-                                               const uint8_t* src_buf,
-                                               FX_DWORD src_size,
-                                               const uint8_t* global_data,
-                                               FX_DWORD global_size,
+                                               CPDF_StreamAcc* src_stream,
+                                               CPDF_StreamAcc* global_stream,
                                                uint8_t* dest_buf,
                                                FX_DWORD dest_pitch,
                                                IFX_Pause* pPause) {
   if (!pJbig2Context) {
     return FXCODEC_STATUS_ERR_PARAMS;
   }
+  JBig2DocumentContext* pJBig2DocumentContext =
+      GetJBig2DocumentContext(this, pPrivateData);
   CCodec_Jbig2Context* m_pJbig2Context = (CCodec_Jbig2Context*)pJbig2Context;
   m_pJbig2Context->m_width = width;
   m_pJbig2Context->m_height = height;
-  m_pJbig2Context->m_src_buf = (unsigned char*)src_buf;
-  m_pJbig2Context->m_src_size = src_size;
-  m_pJbig2Context->m_global_data = global_data;
-  m_pJbig2Context->m_global_size = global_size;
+  m_pJbig2Context->m_pSrcStream = src_stream;
+  m_pJbig2Context->m_pGlobalStream = global_stream;
   m_pJbig2Context->m_dest_buf = dest_buf;
   m_pJbig2Context->m_dest_pitch = dest_pitch;
   m_pJbig2Context->m_pPause = pPause;
   FXSYS_memset(dest_buf, 0, height * dest_pitch);
   m_pJbig2Context->m_pContext = CJBig2_Context::CreateContext(
-      global_data, global_size, src_buf, src_size, &m_SymbolDictCache, pPause);
+      global_stream, src_stream, pJBig2DocumentContext->GetSymbolDictCache(),
+      pPause);
   if (!m_pJbig2Context->m_pContext) {
     return FXCODEC_STATUS_ERROR;
   }
diff --git a/core/src/fxcodec/jbig2/JBig2_BitStream.cpp b/core/src/fxcodec/jbig2/JBig2_BitStream.cpp
index 4a4ed992c7..f39a396e71 100644
--- a/core/src/fxcodec/jbig2/JBig2_BitStream.cpp
+++ b/core/src/fxcodec/jbig2/JBig2_BitStream.cpp
@@ -4,12 +4,18 @@
 
 // Original code copyright 2014 Foxit Software Inc. http://www.foxitsoftware.com
 
+#include "../../../include/fpdfapi/fpdf_objects.h"
 #include "JBig2_BitStream.h"
 
 #include <algorithm>
 
-CJBig2_BitStream::CJBig2_BitStream(const uint8_t* pBuffer, FX_DWORD dwLength)
-    : m_pBuf(pBuffer), m_dwLength(dwLength), m_dwByteIdx(0), m_dwBitIdx(0) {
+CJBig2_BitStream::CJBig2_BitStream(CPDF_StreamAcc* pSrcStream)
+    : m_pBuf(pSrcStream->GetData()),
+      m_dwLength(pSrcStream->GetSize()),
+      m_dwByteIdx(0),
+      m_dwBitIdx(0),
+      m_dwObjNum(pSrcStream->GetStream() ? pSrcStream->GetStream()->GetObjNum()
+                                         : 0) {
   if (m_dwLength > 256 * 1024 * 1024) {
     m_dwLength = 0;
     m_pBuf = nullptr;
@@ -176,3 +182,7 @@ bool CJBig2_BitStream::IsInBound() const {
 FX_DWORD CJBig2_BitStream::LengthInBits() const {
   return m_dwLength << 3;
 }
+
+FX_DWORD CJBig2_BitStream::getObjNum() const {
+  return m_dwObjNum;
+}
diff --git a/core/src/fxcodec/jbig2/JBig2_BitStream.h b/core/src/fxcodec/jbig2/JBig2_BitStream.h
index 9a3d8b0a18..c7c50b8c81 100644
--- a/core/src/fxcodec/jbig2/JBig2_BitStream.h
+++ b/core/src/fxcodec/jbig2/JBig2_BitStream.h
@@ -9,9 +9,11 @@
 
 #include "../../../include/fxcrt/fx_basic.h"
 
+class CPDF_StreamAcc;
+
 class CJBig2_BitStream {
  public:
-  CJBig2_BitStream(const uint8_t* pBuffer, FX_DWORD dwLength);
+  explicit CJBig2_BitStream(CPDF_StreamAcc* pSrcStream);
   ~CJBig2_BitStream();
 
   // TODO(thestig): readFoo() should return bool.
@@ -36,6 +38,7 @@ class CJBig2_BitStream {
   const uint8_t* getPointer() const;
   void offset(FX_DWORD dwOffset);
   FX_DWORD getByteLeft() const;
+  FX_DWORD getObjNum() const;
 
  private:
   void AdvanceBit();
@@ -46,6 +49,7 @@ class CJBig2_BitStream {
   FX_DWORD m_dwLength;
   FX_DWORD m_dwByteIdx;
   FX_DWORD m_dwBitIdx;
+  const FX_DWORD m_dwObjNum;
 
   CJBig2_BitStream(const CJBig2_BitStream&) = delete;
   void operator=(const CJBig2_BitStream&) = delete;
diff --git a/core/src/fxcodec/jbig2/JBig2_Context.cpp b/core/src/fxcodec/jbig2/JBig2_Context.cpp
index 7b11ca2e5e..bb01cb6be8 100644
--- a/core/src/fxcodec/jbig2/JBig2_Context.cpp
+++ b/core/src/fxcodec/jbig2/JBig2_Context.cpp
@@ -35,35 +35,26 @@ size_t GetRefAggContextSize(FX_BOOL val) {
 // list keeps track of the freshness of entries, with freshest ones
 // at the front. Even a tiny cache size like 2 makes a dramatic
 // difference for typical JBIG2 documents.
-//
-// Disabled until we can figure out how to clear cache between documents.
-// https://code.google.com/p/pdfium/issues/detail?id=207
-#define DISABLE_SYMBOL_CACHE
-#ifndef DISABLE_SYMBOL_CACHE
 static const int kSymbolDictCacheMaxSize = 2;
-#endif
 
 CJBig2_Context* CJBig2_Context::CreateContext(
-    const uint8_t* pGlobalData,
-    FX_DWORD dwGlobalLength,
-    const uint8_t* pData,
-    FX_DWORD dwLength,
+    CPDF_StreamAcc* pGlobalStream,
+    CPDF_StreamAcc* pSrcStream,
     std::list<CJBig2_CachePair>* pSymbolDictCache,
     IFX_Pause* pPause) {
-  return new CJBig2_Context(pGlobalData, dwGlobalLength, pData, dwLength,
-                            pSymbolDictCache, pPause);
+  return new CJBig2_Context(pGlobalStream, pSrcStream, pSymbolDictCache, pPause,
+                            false);
 }
 
 void CJBig2_Context::DestroyContext(CJBig2_Context* pContext) {
   delete pContext;
 }
 
-CJBig2_Context::CJBig2_Context(const uint8_t* pGlobalData,
-                               FX_DWORD dwGlobalLength,
-                               const uint8_t* pData,
-                               FX_DWORD dwLength,
+CJBig2_Context::CJBig2_Context(CPDF_StreamAcc* pGlobalStream,
+                               CPDF_StreamAcc* pSrcStream,
                                std::list<CJBig2_CachePair>* pSymbolDictCache,
-                               IFX_Pause* pPause)
+                               IFX_Pause* pPause,
+                               bool bIsGlobal)
     : m_nSegmentDecoded(0),
       m_bInPage(false),
       m_bBufSpecified(false),
@@ -72,15 +63,16 @@ CJBig2_Context::CJBig2_Context(const uint8_t* pGlobalData,
       m_ProcessingStatus(FXCODEC_STATUS_FRAME_READY),
       m_gbContext(NULL),
       m_dwOffset(0),
-      m_pSymbolDictCache(pSymbolDictCache) {
-  if (pGlobalData && (dwGlobalLength > 0)) {
-    m_pGlobalContext = new CJBig2_Context(
-        nullptr, 0, pGlobalData, dwGlobalLength, pSymbolDictCache, pPause);
+      m_pSymbolDictCache(pSymbolDictCache),
+      m_bIsGlobal(bIsGlobal) {
+  if (pGlobalStream && (pGlobalStream->GetSize() > 0)) {
+    m_pGlobalContext = new CJBig2_Context(nullptr, pGlobalStream,
+                                          pSymbolDictCache, pPause, true);
   } else {
     m_pGlobalContext = nullptr;
   }
 
-  m_pStream.reset(new CJBig2_BitStream(pData, dwLength));
+  m_pStream.reset(new CJBig2_BitStream(pSrcStream));
 }
 
 CJBig2_Context::~CJBig2_Context() {
@@ -331,7 +323,8 @@ int32_t CJBig2_Context::parseSegmentHeader(CJBig2_Segment* pSegment) {
   if (m_pStream->readInteger(&pSegment->m_dwData_length) != 0)
     return JBIG2_ERROR_TOO_SHORT;
 
-  pSegment->m_pData = m_pStream->getPointer();
+  pSegment->m_dwObjNum = m_pStream->getObjNum();
+  pSegment->m_dwDataOffset = m_pStream->getOffset();
   pSegment->m_State = JBIG2_SEGMENT_DATA_UNPARSED;
   return JBIG2_SUCCESS;
 }
@@ -602,18 +595,21 @@ int32_t CJBig2_Context::parseSymbolDict(CJBig2_Segment* pSegment,
       JBIG2_memset(grContext.get(), 0, sizeof(JBig2ArithCtx) * size);
     }
   }
-  const uint8_t* key = pSegment->m_pData;
+  CJBig2_CacheKey key =
+      CJBig2_CacheKey(pSegment->m_dwObjNum, pSegment->m_dwDataOffset);
   FX_BOOL cache_hit = false;
   pSegment->m_nResultType = JBIG2_SYMBOL_DICT_POINTER;
-  for (std::list<CJBig2_CachePair>::iterator it = m_pSymbolDictCache->begin();
-       it != m_pSymbolDictCache->end(); ++it) {
-    if (it->first == key) {
-      nonstd::unique_ptr<CJBig2_SymbolDict> copy(it->second->DeepCopy());
-      pSegment->m_Result.sd = copy.release();
-      m_pSymbolDictCache->push_front(*it);
-      m_pSymbolDictCache->erase(it);
-      cache_hit = true;
-      break;
+  if (m_bIsGlobal && key.first != 0) {
+    for (auto it = m_pSymbolDictCache->begin(); it != m_pSymbolDictCache->end();
+         ++it) {
+      if (it->first == key) {
+        nonstd::unique_ptr<CJBig2_SymbolDict> copy(it->second->DeepCopy());
+        pSegment->m_Result.sd = copy.release();
+        m_pSymbolDictCache->push_front(*it);
+        m_pSymbolDictCache->erase(it);
+        cache_hit = true;
+        break;
+      }
     }
   }
   if (!cache_hit) {
@@ -634,17 +630,17 @@ int32_t CJBig2_Context::parseSymbolDict(CJBig2_Segment* pSegment,
         return JBIG2_ERROR_FATAL;
       m_pStream->alignByte();
     }
-#ifndef DISABLE_SYMBOL_CACHE
-    nonstd::unique_ptr<CJBig2_SymbolDict> value =
-        pSegment->m_Result.sd->DeepCopy();
-    if (value && kSymbolDictCacheMaxSize > 0) {
-      while (m_pSymbolDictCache->size() >= kSymbolDictCacheMaxSize) {
-        delete m_pSymbolDictCache->back().second;
-        m_pSymbolDictCache->pop_back();
+    if (m_bIsGlobal && kSymbolDictCacheMaxSize > 0) {
+      nonstd::unique_ptr<CJBig2_SymbolDict> value =
+          pSegment->m_Result.sd->DeepCopy();
+      if (value) {
+        while (m_pSymbolDictCache->size() >= kSymbolDictCacheMaxSize) {
+          delete m_pSymbolDictCache->back().second;
+          m_pSymbolDictCache->pop_back();
+        }
+        m_pSymbolDictCache->push_front(CJBig2_CachePair(key, value.release()));
       }
-      m_pSymbolDictCache->push_front(CJBig2_CachePair(key, value.release()));
     }
-#endif
   }
   if (wFlags & 0x0200) {
     pSegment->m_Result.sd->m_bContextRetained = TRUE;
diff --git a/core/src/fxcodec/jbig2/JBig2_Context.h b/core/src/fxcodec/jbig2/JBig2_Context.h
index 4fcef810f6..61379fe2ff 100644
--- a/core/src/fxcodec/jbig2/JBig2_Context.h
+++ b/core/src/fxcodec/jbig2/JBig2_Context.h
@@ -11,6 +11,7 @@
 #include <utility>
 
 #include "../../../../third_party/base/nonstd_unique_ptr.h"
+#include "../../../include/fpdfapi/fpdf_objects.h"
 #include "../../../include/fxcodec/fx_codec_def.h"
 #include "JBig2_List.h"
 #include "JBig2_Page.h"
@@ -20,7 +21,10 @@ class CJBig2_ArithDecoder;
 class CJBig2_GRDProc;
 class IFX_Pause;
 
-using CJBig2_CachePair = std::pair<const uint8_t*, CJBig2_SymbolDict*>;
+// Cache is keyed by the ObjNum of a stream and an index within the stream.
+using CJBig2_CacheKey = std::pair<FX_DWORD, FX_DWORD>;
+// NB: CJBig2_SymbolDict* is owned.
+using CJBig2_CachePair = std::pair<CJBig2_CacheKey, CJBig2_SymbolDict*>;
 
 #define JBIG2_SUCCESS 0
 #define JBIG2_FAILED -1
@@ -36,10 +40,8 @@ using CJBig2_CachePair = std::pair<const uint8_t*, CJBig2_SymbolDict*>;
 class CJBig2_Context {
  public:
   static CJBig2_Context* CreateContext(
-      const uint8_t* pGlobalData,
-      FX_DWORD dwGlobalLength,
-      const uint8_t* pData,
-      FX_DWORD dwLength,
+      CPDF_StreamAcc* pGlobalStream,
+      CPDF_StreamAcc* pSrcStream,
       std::list<CJBig2_CachePair>* pSymbolDictCache,
       IFX_Pause* pPause = NULL);
 
@@ -55,12 +57,11 @@ class CJBig2_Context {
   FXCODEC_STATUS GetProcessingStatus() { return m_ProcessingStatus; }
 
  private:
-  CJBig2_Context(const uint8_t* pGlobalData,
-                 FX_DWORD dwGlobalLength,
-                 const uint8_t* pData,
-                 FX_DWORD dwLength,
+  CJBig2_Context(CPDF_StreamAcc* pGlobalStream,
+                 CPDF_StreamAcc* pSrcStream,
                  std::list<CJBig2_CachePair>* pSymbolDictCache,
-                 IFX_Pause* pPause);
+                 IFX_Pause* pPause,
+                 bool bIsGlobal);
 
   ~CJBig2_Context();
 
@@ -126,6 +127,7 @@ class CJBig2_Context {
   FX_DWORD m_dwOffset;
   JBig2RegionInfo m_ri;
   std::list<CJBig2_CachePair>* const m_pSymbolDictCache;
+  bool m_bIsGlobal;
 };
 
 #endif  // CORE_SRC_FXCODEC_JBIG2_JBIG2_CONTEXT_H_
diff --git a/core/src/fxcodec/jbig2/JBig2_Segment.cpp b/core/src/fxcodec/jbig2/JBig2_Segment.cpp
index 7ad55f3515..1b0d4e9303 100644
--- a/core/src/fxcodec/jbig2/JBig2_Segment.cpp
+++ b/core/src/fxcodec/jbig2/JBig2_Segment.cpp
@@ -16,7 +16,8 @@ CJBig2_Segment::CJBig2_Segment() {
   m_dwPage_association = 0;
   m_dwData_length = 0;
   m_dwHeader_Length = 0;
-  m_pData = NULL;
+  m_dwObjNum = 0;
+  m_dwDataOffset = 0;
   m_State = JBIG2_SEGMENT_HEADER_UNPARSED;
   m_nResultType = JBIG2_VOID_POINTER;
   m_Result.vd = NULL;
diff --git a/core/src/fxcodec/jbig2/JBig2_Segment.h b/core/src/fxcodec/jbig2/JBig2_Segment.h
index b80a5d282f..4aef27289a 100644
--- a/core/src/fxcodec/jbig2/JBig2_Segment.h
+++ b/core/src/fxcodec/jbig2/JBig2_Segment.h
@@ -50,7 +50,8 @@ class CJBig2_Segment {
   FX_DWORD m_dwData_length;
 
   FX_DWORD m_dwHeader_Length;
-  const uint8_t* m_pData;
+  FX_DWORD m_dwObjNum;
+  FX_DWORD m_dwDataOffset;
   JBig2_SegmentState m_State;
   JBig2_ResultType m_nResultType;
   union {