8 files changed, 278 insertions, 66 deletions

diff --git a/include/mupdf/fitz/context.h b/include/mupdf/fitz/context.h
index 50e8eeaa..47aa8fa5 100644
--- a/include/mupdf/fitz/context.h
+++ b/include/mupdf/fitz/context.h
@@ -361,8 +361,8 @@ struct fz_locks_context_s
 };
 
 enum {
-	FZ_LOCK_ALLOC = 0,
-	FZ_LOCK_FILE, /* Unused now */
+	FZ_LOCK_REAP = 0,
+	FZ_LOCK_ALLOC,
 	FZ_LOCK_FREETYPE,
 	FZ_LOCK_GLYPHCACHE,
 	FZ_LOCK_MAX
diff --git a/include/mupdf/fitz/store.h b/include/mupdf/fitz/store.h
index 2ef9810b..a3c45b7f 100644
--- a/include/mupdf/fitz/store.h
+++ b/include/mupdf/fitz/store.h
@@ -37,7 +37,8 @@ struct fz_storable_s {
 
 struct fz_key_storable_s {
 	fz_storable storable;
-	int store_key_refs;
+	short store_key_refs;
+	unsigned short needs_reaping;
 };
 
 #define FZ_INIT_STORABLE(S_,RC,DROP) \
@@ -47,14 +48,14 @@ struct fz_key_storable_s {
 
 #define FZ_INIT_KEY_STORABLE(KS_,RC,DROP) \
 	do { fz_key_storable *KS = &(KS_)->key_storable; KS->store_key_refs = 0;\
-	FZ_INIT_STORABLE(KS,RC,DROP); \
+	KS->needs_reaping = 0; FZ_INIT_STORABLE(KS,RC,DROP); \
 	} while (0)
 
 void *fz_keep_storable(fz_context *, const fz_storable *);
 void fz_drop_storable(fz_context *, const fz_storable *);
 
 void *fz_keep_key_storable(fz_context *, const fz_key_storable *);
-int fz_drop_key_storable(fz_context *, const fz_key_storable *);
+void fz_drop_key_storable(fz_context *, const fz_key_storable *);
 
 void *fz_keep_key_storable_key(fz_context *, const fz_key_storable *);
 void fz_drop_key_storable_key(fz_context *, const fz_key_storable *);
@@ -74,6 +75,39 @@ void fz_drop_key_storable_key(fz_context *, const fz_key_storable *);
 	to an fz_store_hash structure. If make_hash_key function returns 0,
 	then the key is determined not to be hashable, and the value is
 	not stored in the hash table.
+
+	Some objects can be used both as values within the store, and as a
+	component of keys within the store. We refer to these objects as
+	"key storable" objects. In this case, we need to take additional
+	care to ensure that we do not end up keeping an item within the
+	store, purely because it's value is referred to by another key in
+	the store.
+
+	An example of this are fz_images in PDF files. Each fz_image is
+	placed into the	store to enable it to be easily reused. When the
+	image is rendered, a pixmap is generated from the image, and the
+	pixmap is placed into the store so it can be reused on subsequent
+	renders. The image forms part of the key for the pixmap.
+
+	When we close the pdf document (and any associated pages/display
+	lists etc), we drop the images from the store. This may leave us
+	in the position of the images having non-zero reference counts
+	purely because they are used as part of the keys for the pixmaps.
+
+	We therefore use special reference counting functions to keep
+	track of these "key storable" items, and hence store the number of
+	references to these items that are used in keys.
+
+	When the number of references to an object == the number of
+	references to an object from keys in the store, we know that we can
+	remove all the items which have that object as part of the key.
+	This is done by running a pass over the store, 'reaping' those
+	items.
+
+	Reap passes are slower than we would like as they touching every
+	item in the store. We therefore provide a way to 'batch' such
+	reap passes together, using fz_defer_reap_start/fz_defer_reap_end
+	to bracket a region in which many may be triggered.
 */
 typedef struct fz_store_hash_s fz_store_hash;
 
@@ -110,6 +144,7 @@ struct fz_store_type_s
 	void (*drop_key)(fz_context *,void *);
 	int (*cmp_key)(fz_context *ctx, void *, void *);
 	void (*print)(fz_context *ctx, fz_output *out, void *);
+	int (*needs_reap)(fz_context *ctx, void *);
 };
 
 /*
@@ -219,4 +254,33 @@ void fz_filter_store(fz_context *ctx, fz_store_filter_fn *fn, void *arg, fz_stor
 void fz_print_store(fz_context *ctx, fz_output *out);
 void fz_print_store_locked(fz_context *ctx, fz_output *out);
 
+/*
+	fz_defer_reap_start: Increment the defer reap count.
+
+	No reap operations will take place (except for those
+	triggered by an immediate failed malloc) until the
+	defer reap count returns to 0.
+
+	Call this at the start of a process during which you
+	potentially might drop many reapable objects.
+
+	It is vital that every fz_defer_reap_start is matched
+	by a fz_defer_reap_end call.
+*/
+void fz_defer_reap_start(fz_context *ctx);
+
+/*
+	fz_defer_reap_end: Decrement the defer reap count.
+
+	If the defer reap count returns to 0, and the store
+	has reapable objects in, a reap pass will begin.
+
+	Call this at the end of a process during which you
+	potentially might drop many reapable objects.
+
+	It is vital that every fz_defer_reap_start is matched
+	by a fz_defer_reap_end call.
+*/
+void fz_defer_reap_end(fz_context *ctx);
+
 #endif
diff --git a/source/fitz/image.c b/source/fitz/image.c
index 247cb088..8d692fbd 100644
--- a/source/fitz/image.c
+++ b/source/fitz/image.c
@@ -82,39 +82,34 @@ fz_cmp_image_key(fz_context *ctx, void *k0_, void *k1_)
 }
 
 static void
-fz_print_image(fz_context *ctx, fz_output *out, void *key_)
+fz_print_image_key(fz_context *ctx, fz_output *out, void *key_)
 {
 	fz_image_key *key = (fz_image_key *)key_;
 	fz_printf(ctx, out, "(image %d x %d sf=%d) ", key->image->w, key->image->h, key->l2factor);
 }
 
+static int
+fz_needs_reap_image_key(fz_context *ctx, void *key_)
+{
+	fz_image_key *key = (fz_image_key *)key_;
+
+	return (key->image->key_storable.needs_reaping);
+}
+
 static fz_store_type fz_image_store_type =
 {
 	fz_make_hash_image_key,
 	fz_keep_image_key,
 	fz_drop_image_key,
 	fz_cmp_image_key,
-	fz_print_image
+	fz_print_image_key,
+	fz_needs_reap_image_key
 };
 
-static int
-drop_matching_images(fz_context *ctx, void *image_, void *key_)
-{
-	fz_image_key *key = (fz_image_key *)key_;
-	fz_image *image = (fz_image *)image_;
-
-	return key->image == image;
-}
-
 void
 fz_drop_image(fz_context *ctx, fz_image *image)
 {
-	if (fz_drop_key_storable(ctx, &image->key_storable))
-	{
-		/* All the image refs left are references from keys in the store. */
-		/* We can never hope to match these keys again, so drop the objects. */
-		fz_filter_store(ctx, drop_matching_images, image, &fz_image_store_type);
-	}
+	fz_drop_key_storable(ctx, &image->key_storable);
 }
 
 static void
diff --git a/source/fitz/list-device.c b/source/fitz/list-device.c
index dc8b1f0b..3d17ee7d 100644
--- a/source/fitz/list-device.c
+++ b/source/fitz/list-device.c
@@ -1369,7 +1369,9 @@ fz_keep_display_list(fz_context *ctx, fz_display_list *list)
 void
 fz_drop_display_list(fz_context *ctx, fz_display_list *list)
 {
+	fz_defer_reap_start(ctx);
 	fz_drop_storable(ctx, &list->storable);
+	fz_defer_reap_end(ctx);
 }
 
 fz_rect *
diff --git a/source/fitz/store.c b/source/fitz/store.c
index d9674eb6..639e0662 100644
--- a/source/fitz/store.c
+++ b/source/fitz/store.c
@@ -29,6 +29,10 @@ struct fz_store_s
 	/* We keep track of the size of the store, and keep it below max. */
 	size_t max;
 	size_t size;
+
+	/* Protected by the reap lock */
+	int defer_reap_count;
+	int needs_reaping;
 };
 
 void
@@ -50,6 +54,8 @@ fz_new_store_context(fz_context *ctx, size_t max)
 	store->tail = NULL;
 	store->size = 0;
 	store->max = max;
+	store->defer_reap_count = 0;
+	store->needs_reaping = 0;
 	ctx->store = store;
 }
 
@@ -86,16 +92,93 @@ void *fz_keep_key_storable(fz_context *ctx, const fz_key_storable *sc)
 	return fz_keep_storable(ctx, &sc->storable);
 }
 
-int fz_drop_key_storable(fz_context *ctx, const fz_key_storable *sc)
+/*
+	Entered with FZ_LOCK_ALLOC and FZ_LOCK_REAP held.
+	Drops FZ_LOCK_ALLOC.
+*/
+static void
+do_reap(fz_context *ctx)
+{
+	fz_store *store = ctx->store;
+	fz_item *item, *prev, *remove;
+
+	if (store == NULL)
+	{
+		fz_unlock(ctx, FZ_LOCK_ALLOC);
+		return;
+	}
+
+	fz_assert_lock_held(ctx, FZ_LOCK_ALLOC);
+	fz_assert_lock_held(ctx, FZ_LOCK_REAP);
+
+	ctx->store->needs_reaping = 0;
+
+	/* Reap the items */
+	remove = NULL;
+	for (item = store->tail; item; item = prev)
+	{
+		prev = item->prev;
+
+		if (item->type->needs_reap == NULL || item->type->needs_reap(ctx, item->key) == 0)
+			continue;
+
+		/* We have to drop it */
+		store->size -= item->size;
+
+		/* Unlink from the linked list */
+		if (item->next)
+			item->next->prev = item->prev;
+		else
+			store->tail = item->prev;
+		if (item->prev)
+			item->prev->next = item->next;
+		else
+			store->head = item->next;
+
+		/* Remove from the hash table */
+		if (item->type->make_hash_key)
+		{
+			fz_store_hash hash = { NULL };
+			hash.drop = item->val->drop;
+			if (item->type->make_hash_key(ctx, &hash, item->key))
+				fz_hash_remove(ctx, store->hash, &hash);
+		}
+
+		/* Store whether to drop this value or not in 'prev' */
+		item->prev = (item->val->refs > 0 && --item->val->refs == 0) ? item : NULL;
+
+		/* Store it in our removal chain - just singly linked */
+		item->next = remove;
+		remove = item;
+	}
+	fz_unlock(ctx, FZ_LOCK_ALLOC);
+
+	/* Now drop the remove chain */
+	for (item = remove; item != NULL; item = remove)
+	{
+		remove = item->next;
+
+		/* Drop a reference to the value (freeing if required) */
+		if (item->prev)
+			item->val->drop(ctx, item->val);
+
+		/* Always drops the key and drop the item */
+		item->type->drop_key(ctx, item->key);
+		fz_free(ctx, item);
+	}
+
+}
+
+void fz_drop_key_storable(fz_context *ctx, const fz_key_storable *sc)
 {
 	/* Explicitly drop const to allow us to use const
 	 * sanely throughout the code. */
 	fz_key_storable *s = (fz_key_storable *)sc;
 	int drop;
-	int only_refs_are_store_key_refs = 0;
+	int unlock = 1;
 
 	if (s == NULL)
-		return 0;
+		return;
 
 	if (s->storable.refs > 0)
 		(void)Memento_dropRef(s);
@@ -104,11 +187,22 @@ int fz_drop_key_storable(fz_context *ctx, const fz_key_storable *sc)
 	{
 		drop = --s->storable.refs == 0;
 		if (!drop && s->storable.refs == s->store_key_refs)
-			only_refs_are_store_key_refs = 1;
+		{
+			fz_lock(ctx, FZ_LOCK_REAP);
+			if (ctx->store->defer_reap_count > 0)
+				ctx->store->needs_reaping = 1;
+			else
+			{
+				do_reap(ctx);
+				unlock = 0;
+			}
+			fz_unlock(ctx, FZ_LOCK_REAP);
+		}
 	}
 	else
 		drop = 0;
-	fz_unlock(ctx, FZ_LOCK_ALLOC);
+	if (unlock)
+		fz_unlock(ctx, FZ_LOCK_ALLOC);
 	/*
 		If we are dropping the last reference to an object, then
 		it cannot possibly be in the store (as the store always
@@ -118,7 +212,6 @@ int fz_drop_key_storable(fz_context *ctx, const fz_key_storable *sc)
 	 */
 	if (drop)
 		s->storable.drop(ctx, &s->storable);
-	return only_refs_are_store_key_refs;
 }
 
 void *fz_keep_key_storable_key(fz_context *ctx, const fz_key_storable *sc)
@@ -385,8 +478,25 @@ fz_store_item(fz_context *ctx, void *key, void *val_, size_t itemsize, fz_store_
 		size = store->size + itemsize;
 		while (size > store->max)
 		{
+			size_t saved;
+			int relock = 0;
+
+			/* First, do any outstanding reaping, even if defer_reap_count > 0 */
+			fz_lock(ctx, FZ_LOCK_REAP);
+			if (store->needs_reaping)
+			{
+				do_reap(ctx); /* Drops alloc lock */
+				relock = 1;
+			}
+			fz_unlock(ctx, FZ_LOCK_REAP);
+			if (relock)
+				fz_lock(ctx, FZ_LOCK_ALLOC);
+			size = store->size + itemsize;
+			if (size <= store->max)
+				break;
+
 			/* ensure_space may drop, then retake the lock */
-			size_t saved = ensure_space(ctx, size - store->max);
+			saved = ensure_space(ctx, size - store->max);
 			size -= saved;
 			if (saved == 0)
 			{
@@ -794,5 +904,32 @@ void fz_filter_store(fz_context *ctx, fz_store_filter_fn *fn, void *arg, fz_stor
 		item->type->drop_key(ctx, item->key);
 		fz_free(ctx, item);
 	}
+}
+
+void fz_defer_reap_start(fz_context *ctx)
+{
+	if (ctx == NULL || ctx->store == NULL)
+		return;
+
+	fz_lock(ctx, FZ_LOCK_REAP);
+	ctx->store->defer_reap_count++;
+	fz_unlock(ctx, FZ_LOCK_REAP);
+}
+
+void fz_defer_reap_end(fz_context *ctx)
+{
+	int reap;
+
+	if (ctx == NULL || ctx->store == NULL)
+		return;
 
+	fz_lock(ctx, FZ_LOCK_ALLOC);
+	fz_lock(ctx, FZ_LOCK_REAP);
+	--ctx->store->defer_reap_count;
+	reap = ctx->store->defer_reap_count == 0 && ctx->store->needs_reaping;
+	if (reap)
+		do_reap(ctx); /* Drops FZ_LOCK_ALLOC */
+	fz_unlock(ctx, FZ_LOCK_REAP);
+	if (!reap)
+		fz_unlock(ctx, FZ_LOCK_ALLOC);
 }
diff --git a/source/pdf/pdf-interpret.c b/source/pdf/pdf-interpret.c
index f3840794..00c4176c 100644
--- a/source/pdf/pdf-interpret.c
+++ b/source/pdf/pdf-interpret.c
@@ -1231,12 +1231,14 @@ pdf_process_contents(fz_context *ctx, pdf_processor *proc, pdf_document *doc, pd
 
 	fz_try(ctx)
 	{
+		fz_defer_reap_start(ctx);
 		stm = pdf_open_contents_stream(ctx, doc, stmobj);
 		pdf_process_stream(ctx, proc, &csi, stm);
 		pdf_process_end(ctx, proc, &csi);
 	}
 	fz_always(ctx)
 	{
+		fz_defer_reap_end(ctx);
 		fz_drop_stream(ctx, stm);
 		pdf_clear_stack(ctx, &csi);
 		pdf_lexbuf_fin(ctx, &buf);
diff --git a/source/pdf/pdf-store.c b/source/pdf/pdf-store.c
index 44575691..33d34ab9 100644
--- a/source/pdf/pdf-store.c
+++ b/source/pdf/pdf-store.c
@@ -47,7 +47,8 @@ static fz_store_type pdf_obj_store_type =
 	pdf_keep_key,
 	pdf_drop_key,
 	pdf_cmp_key,
-	pdf_print_key
+	pdf_print_key,
+	NULL
 };
 
 void
diff --git a/source/pdf/pdf-xref.c b/source/pdf/pdf-xref.c
index c4db8e4f..7d21775a 100644
--- a/source/pdf/pdf-xref.c
+++ b/source/pdf/pdf-xref.c
@@ -1569,54 +1569,65 @@ pdf_drop_document_imp(fz_context *ctx, pdf_document *doc)
 	if (!doc)
 		return;
 
-	/* Type3 glyphs in the glyph cache can contain pdf_obj pointers
-	 * that we are about to destroy. Simplest solution is to bin the
-	 * glyph cache at this point. */
-	fz_purge_glyph_cache(ctx);
+	fz_try(ctx)
+	{
+		fz_defer_reap_start(ctx);
 
-	if (doc->js)
-		pdf_drop_js(ctx, doc->js);
+		/* Type3 glyphs in the glyph cache can contain pdf_obj pointers
+		 * that we are about to destroy. Simplest solution is to bin the
+		 * glyph cache at this point. */
+		fz_purge_glyph_cache(ctx);
 
-	pdf_drop_xref_sections(ctx, doc);
-	fz_free(ctx, doc->xref_index);
+		if (doc->js)
+			pdf_drop_js(ctx, doc->js);
 
-	if (doc->focus_obj)
-		pdf_drop_obj(ctx, doc->focus_obj);
-	if (doc->file)
-		fz_drop_stream(ctx, doc->file);
-	if (doc->crypt)
-		pdf_drop_crypt(ctx, doc->crypt);
+		pdf_drop_xref_sections(ctx, doc);
+		fz_free(ctx, doc->xref_index);
 
-	pdf_drop_obj(ctx, doc->linear_obj);
-	if (doc->linear_page_refs)
-	{
-		for (i=0; i < doc->page_count; i++)
+		if (doc->focus_obj)
+			pdf_drop_obj(ctx, doc->focus_obj);
+		if (doc->file)
+			fz_drop_stream(ctx, doc->file);
+		if (doc->crypt)
+			pdf_drop_crypt(ctx, doc->crypt);
+
+		pdf_drop_obj(ctx, doc->linear_obj);
+		if (doc->linear_page_refs)
 		{
-			pdf_drop_obj(ctx, doc->linear_page_refs[i]);
+			for (i=0; i < doc->page_count; i++)
+			{
+				pdf_drop_obj(ctx, doc->linear_page_refs[i]);
+			}
+			fz_free(ctx, doc->linear_page_refs);
 		}
-		fz_free(ctx, doc->linear_page_refs);
-	}
-	fz_free(ctx, doc->hint_page);
-	fz_free(ctx, doc->hint_shared_ref);
-	fz_free(ctx, doc->hint_shared);
-	fz_free(ctx, doc->hint_obj_offsets);
+		fz_free(ctx, doc->hint_page);
+		fz_free(ctx, doc->hint_shared_ref);
+		fz_free(ctx, doc->hint_shared);
+		fz_free(ctx, doc->hint_obj_offsets);
 
-	for (i=0; i < doc->num_type3_fonts; i++)
-	{
-		fz_decouple_type3_font(ctx, doc->type3_fonts[i], (void *)doc);
-		fz_drop_font(ctx, doc->type3_fonts[i]);
-	}
-	fz_free(ctx, doc->type3_fonts);
+		for (i=0; i < doc->num_type3_fonts; i++)
+		{
+			fz_decouple_type3_font(ctx, doc->type3_fonts[i], (void *)doc);
+			fz_drop_font(ctx, doc->type3_fonts[i]);
+		}
+		fz_free(ctx, doc->type3_fonts);
 
-	pdf_drop_ocg(ctx, doc->ocg);
+		pdf_drop_ocg(ctx, doc->ocg);
 
-	pdf_empty_store(ctx, doc);
+		pdf_empty_store(ctx, doc);
 
-	pdf_lexbuf_fin(ctx, &doc->lexbuf.base);
+		pdf_lexbuf_fin(ctx, &doc->lexbuf.base);
 
-	pdf_drop_resource_tables(ctx, doc);
+		pdf_drop_resource_tables(ctx, doc);
 
-	fz_free(ctx, doc);
+		fz_free(ctx, doc);
+	}
+	fz_always(ctx)
+	{
+		fz_defer_reap_end(ctx);
+	}
+	fz_catch(ctx)
+		fz_rethrow(ctx);
 }
 
 void