Optimize the SkRawStream when the input is an asset stream

src/codec/SkRawCodec.cpp

 /*
  * Copyright 2016 Google Inc.
  *
  * Use of this source code is governed by a BSD-style license that can be
  * found in the LICENSE file.
  */
 
 #include "SkCodec.h"
 #include "SkCodecPriv.h"
 #include "SkColorPriv.h"
@@ skipping 139 matching lines @@
         // To avoid arbitary allocation requests which might lead to out-of-memory, limit the
         // amount of memory that can be allocated at once. The memory limit is based on experiments
         // and supposed to be sufficient for all valid DNG images.
         if (size > 300 * 1024 * 1024) {  // 300 MB
             ThrowMemoryFull();
         }
         return dng_memory_allocator::Allocate(size);
     }
 };
 
+bool is_asset_stream(const SkStream& stream) {
+    return stream.hasLength() && stream.hasPosition();
+}
+
 }  // namespace
 
-// Note: this class could throw exception if it is used as dng_stream.
-class SkRawStream : public ::piex::StreamInterface {
+class SkRawStream {
 public:
-    // Note that this call will take the ownership of stream.
-    explicit SkRawStream(SkStream* stream)
-        : fStream(stream), fWholeStreamRead(false) {}
+    virtual ~SkRawStream() {}
 
-    ~SkRawStream() override {}
+   /* 
+    * Gets the length of the stream. Depending on the type of stream, this may require reading to
+    * the end of the stream.
+    */
+   virtual uint64 getLength() = 0;
+
+   virtual bool read(void* data, size_t offset, size_t length) = 0;
 
     /*
      * Creates an SkMemoryStream from the offset with size.
      * Note: for performance reason, this function is destructive to the SkRawStream. One should
      *       abandon current object after the function call.
      */
-    SkMemoryStream* transferBuffer(size_t offset, size_t size) {
+   virtual SkMemoryStream* transferBuffer(size_t offset, size_t size) = 0;
+};
+
+class SkRawBufferedStream : public SkRawStream {
+public:
+    // Will take the ownership of the stream.
+    explicit SkRawBufferedStream(SkStream* stream)
+        : fStream(stream)
+        , fWholeStreamRead(false)
+    {
+        // Only use SkRawBufferedStream when the stream is not an asset stream.
+        SkASSERT(!is_asset_stream(*stream));
+    }
+
+    ~SkRawBufferedStream() override {}
+
+    uint64 getLength() override {
+        if (!this->bufferMoreData(kReadToEnd)) {  // read whole stream
+            ThrowReadFile();
+        }
+        return fStreamBuffer.bytesWritten();
+    }
+
+    bool read(void* data, size_t offset, size_t length) override {
+        if (length == 0) {
+            return true;
+        }
+
+        size_t sum;
+        if (!safe_add_to_size_t(offset, length, &sum)) {
+            return false;
+        }
+
+        return this->bufferMoreData(sum) && fStreamBuffer.read(data, offset, length);
+    }
+
+    SkMemoryStream* transferBuffer(size_t offset, size_t size) override {
         SkAutoTUnref<SkData> data(SkData::NewUninitialized(size));
         if (offset > fStreamBuffer.bytesWritten()) {
             // If the offset is not buffered, read from fStream directly and skip the buffering.
             const size_t skipLength = offset - fStreamBuffer.bytesWritten();
             if (fStream->skip(skipLength) != skipLength) {
                 return nullptr;
             }
             const size_t bytesRead = fStream->read(data->writable_data(), size);
             if (bytesRead < size) {
                 data.reset(SkData::NewSubset(data.get(), 0, bytesRead));
@@ skipping 14 matching lines @@
                     if (!safe_add_to_size_t(alreadyBuffered, bytesRead, &newSize)) {
                         return nullptr;
                     }
                     data.reset(SkData::NewSubset(data.get(), 0, newSize));
                 }
             }
         }
         return new SkMemoryStream(data);
     }
 
-    // For PIEX
-    ::piex::Error GetData(const size_t offset, const size_t length,
-                          uint8* data) override {
-        if (offset == 0 && length == 0) {
-            return ::piex::Error::kOk;
-        }
-        size_t sum;
-        if (!safe_add_to_size_t(offset, length, &sum) || !this->bufferMoreData(sum)) {
-            return ::piex::Error::kFail;
-        }
-        if (!fStreamBuffer.read(data, offset, length)) {
-            return ::piex::Error::kFail;
-        }
-        return ::piex::Error::kOk;
-    }
-
-    // For dng_stream
-    uint64 getLength() {
-        if (!this->bufferMoreData(kReadToEnd)) {  // read whole stream
-            ThrowReadFile();
-        }
-        return fStreamBuffer.bytesWritten();
-    }
-
-    // For dng_stream
-    void read(void* data, uint32 count, uint64 offset) {
-        if (count == 0 && offset == 0) {
-            return;
-        }
-        size_t sum;
-        if (!safe_add_to_size_t(static_cast<uint64>(count), offset, &sum) ||
-            !this->bufferMoreData(sum)) {
-            ThrowReadFile();
-        }
-
-        if (!fStreamBuffer.read(data, static_cast<size_t>(offset), count)) {
-            ThrowReadFile();
-        }
-    }
-
 private:
     // Note: if the newSize == kReadToEnd (0), this function will read to the end of stream.
     bool bufferMoreData(size_t newSize) {
         if (newSize == kReadToEnd) {
             if (fWholeStreamRead) {  // already read-to-end.
                 return true;
             }
 
             // TODO: optimize for the special case when the input is SkMemoryStream.
             return SkStreamCopy(&fStreamBuffer, fStream.get());
@@ skipping 19 matching lines @@
     }
 
     SkAutoTDelete<SkStream> fStream;
     bool fWholeStreamRead;
 
     SkDynamicMemoryWStream fStreamBuffer;
 
     const size_t kReadToEnd = 0;
 };
 
+class SkRawAssetStream : public SkRawStream {
+public:
+    // Will take the ownership of the stream.
+    explicit SkRawAssetStream(SkStream* stream)
+        : fStream(stream)
+    {
+        // Only use SkRawAssetStream when the stream is an asset stream.
+        SkASSERT(is_asset_stream(*stream));
+    }
+
+    ~SkRawAssetStream() override {}
+
+    uint64 getLength() override {
+        return fStream->getLength();
+    }
+
+
+    bool read(void* data, size_t offset, size_t length) override {
+        if (length == 0) {
+            return true;
+        }
+
+        size_t sum;
+        if (!safe_add_to_size_t(offset, length, &sum)) {
+            return false;
+        }
+
+        return fStream->seek(offset) && (fStream->read(data, length) == length);
+    }
+
+    SkMemoryStream* transferBuffer(size_t offset, size_t size) override {
+        if (fStream->getLength() < offset) {
+            return nullptr;
+        }
+
+        size_t sum;
+        if (!safe_add_to_size_t(offset, size, &sum)) {
+            return nullptr;
+        }
+
+        // This will allow read less than the requested "size", because the JPEG codec wants to
+        // handle also a partial JPEG file.
+        const size_t bytesToRead = SkTMin(sum, fStream->getLength()) - offset;
+        if (bytesToRead == 0) {
+            return nullptr;
+        }
+
+        if (fStream->getMemoryBase()) {  // directly copy if getMemoryBase() is available.
+            SkAutoTUnref<SkData> data(SkData::NewWithCopy(
+                static_cast<const uint8_t*>(fStream->getMemoryBase()) + offset, bytesToRead));
+            fStream.free();
+            return new SkMemoryStream(data);
+        } else {
+            SkAutoTUnref<SkData> data(SkData::NewUninitialized(bytesToRead));
+            if (!fStream->seek(offset)) {
+                return nullptr;
+            }
+            const size_t bytesRead = fStream->read(data->writable_data(), bytesToRead);
+            if (bytesRead < bytesToRead) {
+                data.reset(SkData::NewSubset(data.get(), 0, bytesRead));
+            }
+            return new SkMemoryStream(data);
+        }
+    }
+private:
+    SkAutoTDelete<SkStream> fStream;
+};
+
+class SkPiexStream : public ::piex::StreamInterface {
+public:
+    // Will NOT take the ownership of the stream.
+    explicit SkPiexStream(SkRawStream* stream) : fStream(stream) {}
+
+    ~SkPiexStream() override {}
+
+    ::piex::Error GetData(const size_t offset, const size_t length,
+                          uint8* data) override {
+        return fStream->read(static_cast<void*>(data), offset, length) ?
+            ::piex::Error::kOk : ::piex::Error::kFail;
+    }
+
+private:
+    SkRawStream* fStream;
+};
+
 class SkDngStream : public dng_stream {
 public:
-    SkDngStream(SkRawStream* rawStream) : fRawStream(rawStream) {}
+    // Will NOT take the ownership of the stream.
+    SkDngStream(SkRawStream* stream) : fStream(stream) {}
 
-    uint64 DoGetLength() override { return fRawStream->getLength(); }
+    ~SkDngStream() override {}
+
+    uint64 DoGetLength() override { return fStream->getLength(); }
 
     void DoRead(void* data, uint32 count, uint64 offset) override {
-        fRawStream->read(data, count, offset);
+        size_t sum;
+        if (!safe_add_to_size_t(static_cast<uint64>(count), offset, &sum) ||
+            !fStream->read(data, static_cast<size_t>(offset), static_cast<size_t>(count))) {
+            ThrowReadFile();
+        }
     }
 
 private:
-    SkRawStream* fRawStream;
+    SkRawStream* fStream;
 };
 
 class SkDngImage {
 public:
+    // Will take the ownership of the stream.
     static SkDngImage* NewFromStream(SkRawStream* stream) {
         SkAutoTDelete<SkDngImage> dngImage(new SkDngImage(stream));
         if (!dngImage->readDng()) {
             return nullptr;
         }
 
         SkASSERT(dngImage->fNegative);
         return dngImage.release();
     }
 
@@ skipping 116 matching lines @@
     bool fIsScalable;
     bool fIsXtransImage;
 };
 
 /*
  * Tries to handle the image with PIEX. If PIEX returns kOk and finds the preview image, create a
  * SkJpegCodec. If PIEX returns kFail, then the file is invalid, return nullptr. In other cases,
  * fallback to create SkRawCodec for DNG images.
  */
 SkCodec* SkRawCodec::NewFromStream(SkStream* stream) {
-    SkAutoTDelete<SkRawStream> rawStream(new SkRawStream(stream));
+    SkAutoTDelete<SkRawStream> rawStream;
+    if (is_asset_stream(*stream)) {
+        rawStream.reset(new SkRawAssetStream(stream));
+    } else {
+        rawStream.reset(new SkRawBufferedStream(stream));
+    }
+
+    // Does not take the ownership of rawStream.
+    SkPiexStream piexStream(rawStream.get());
     ::piex::PreviewImageData imageData;
-    if (::piex::IsRaw(rawStream.get())) {
-        ::piex::Error error = ::piex::GetPreviewImageData(rawStream.get(), &imageData);
+    if (::piex::IsRaw(&piexStream)) {
+        ::piex::Error error = ::piex::GetPreviewImageData(&piexStream, &imageData);
 
         if (error == ::piex::Error::kOk && imageData.preview_length > 0) {
 #if !defined(GOOGLE3)
             // transferBuffer() is destructive to the rawStream. Abandon the rawStream after this
             // function call.
             // FIXME: one may avoid the copy of memoryStream and use the buffered rawStream.
             SkMemoryStream* memoryStream =
-                    rawStream->transferBuffer(imageData.preview_offset, imageData.preview_length);
+                rawStream->transferBuffer(imageData.preview_offset, imageData.preview_length);
             return memoryStream ? SkJpegCodec::NewFromStream(memoryStream) : nullptr;
 #else
             return nullptr;
 #endif
         } else if (error == ::piex::Error::kFail) {
             return nullptr;
         }
     }
 
+    // Takes the ownership of the rawStream.
     SkAutoTDelete<SkDngImage> dngImage(SkDngImage::NewFromStream(rawStream.release()));
     if (!dngImage) {
         return nullptr;
     }
 
     return new SkRawCodec(dngImage.release());
 }
 
 SkCodec::Result SkRawCodec::onGetPixels(const SkImageInfo& requestedInfo, void* dst,
                                         size_t dstRowBytes, const Options& options,
@@ skipping 86 matching lines @@
     SkISize sizeFloor = this->onGetScaledDimensions(1.f / std::floor(fullShortEdge / shortEdge));
     SkISize sizeCeil = this->onGetScaledDimensions(1.f / std::ceil(fullShortEdge / shortEdge));
     return sizeFloor == dim || sizeCeil == dim;
 }
 
 SkRawCodec::~SkRawCodec() {}
 
 SkRawCodec::SkRawCodec(SkDngImage* dngImage)
     : INHERITED(dngImage->getImageInfo(), nullptr)
     , fDngImage(dngImage) {}