--- a/yaffs_checkptrw.c
+++ b/yaffs_checkptrw.c
@@ -13,71 +13,164 @@
  */
 
 const char *yaffs_checkptrw_c_version =
-    "$Id: yaffs_checkptrw.c,v 1.1 2006-05-08 10:13:34 charles Exp $";
+    "$Id: yaffs_checkptrw.c,v 1.2 2006-05-17 09:31:07 charles Exp $";
 
 
 #include "yaffs_checkptrw.h"
 
 
-int yaffs_CheckpointSpaceOk(yaffs_Device *dev)
+static int yaffs_CheckpointSpaceOk(yaffs_Device *dev)
 {
 
-	T(YAFFS_TRACE_CHECKPOINT,(TSTR("checkpt blocks %d %d %d %d" TENDSTR),
-		dev->checkpointStartBlock, dev->checkpointEndBlock,
-		dev->startBlock,dev->endBlock));
-		
-	if(dev->checkpointStartBlock >= dev->checkpointEndBlock)
-		return 0;
+	int blocksAvailable = dev->nErasedBlocks - dev->nReservedBlocks;
+	
+	if(blocksAvailable < 0)
+		blocksAvailable = 0;
 		
-	if(dev->checkpointStartBlock >= dev->startBlock && 
-	   dev->checkpointStartBlock <= dev->endBlock)
-	   	return 0;
-
-	if(dev->checkpointEndBlock >= dev->startBlock && 
-	   dev->checkpointEndBlock <= dev->endBlock)
-	   	return 0;
+	T(YAFFS_TRACE_CHECKPOINT,(TSTR("checkpt blocks available" TENDSTR)));
+		
+	
+	if(blocksAvailable <= 0)
+		return 0;
 	   	
 	return 1;
 }
 
+
+
+static int yaffs_CheckpointErase(yaffs_Device *dev)
+{
+	
+	int i;
+	
+
+	if(!dev->eraseBlockInNAND)	
+		return 0;
+	T(YAFFS_TRACE_CHECKPOINT,(TSTR("checking blocks %d to %d"TENDSTR),
+		dev->startBlock,dev->endBlock));
+		
+	for(i = dev->startBlock; i <= dev->endBlock; i++) {
+		yaffs_BlockInfo *bi = &dev->blockInfo[i];
+		if(bi->blockState == YAFFS_BLOCK_STATE_CHECKPOINT){
+			T(YAFFS_TRACE_CHECKPOINT,(TSTR("erasing checkpt block %d"TENDSTR),i));
+			if(dev->eraseBlockInNAND(dev,i)){
+				bi->blockState = YAFFS_BLOCK_STATE_EMPTY;
+				dev->nErasedBlocks++;
+				dev->nFreeChunks += dev->nChunksPerBlock;
+			}
+			else {
+				dev->markNANDBlockBad(dev,i);
+				bi->blockState = YAFFS_BLOCK_STATE_DEAD;
+			}
+		}
+	}
+	
+	dev->blocksInCheckpoint = 0;
+	
+	return 1;
+}
+
+
+static void yaffs_CheckpointFindNextErasedBlock(yaffs_Device *dev)
+{
+	int  i;
+	int blocksAvailable = dev->nErasedBlocks - dev->nReservedBlocks;
+		
+	if(dev->checkpointNextBlock >= 0 &&
+	   dev->checkpointNextBlock <= dev->endBlock &&
+	   blocksAvailable > 0){
+	
+		for(i = dev->checkpointNextBlock; i <= dev->endBlock; i++){
+			yaffs_BlockInfo *bi = &dev->blockInfo[i];
+			if(bi->blockState == YAFFS_BLOCK_STATE_EMPTY){
+				dev->checkpointNextBlock = i + 1;
+				dev->nErasedBlocks--;
+				dev->checkpointCurrentBlock = i;
+				T(YAFFS_TRACE_CHECKPOINT,(TSTR("allocating checkpt block %d"TENDSTR),i));
+				return;
+			}
+		}
+	}
+	T(YAFFS_TRACE_CHECKPOINT,(TSTR("out of checkpt blocks"TENDSTR)));
+	
+	dev->checkpointNextBlock = -1;
+	dev->checkpointCurrentBlock = -1;
+}
+
+static void yaffs_CheckpointFindNextCheckpointBlock(yaffs_Device *dev)
+{
+	int  i;
+	yaffs_ExtendedTags tags;
+	
+	if(dev->blocksInCheckpoint < dev->checkpointMaxBlocks) 
+		for(i = dev->checkpointNextBlock; i <= dev->endBlock; i++){
+			int chunk = i * dev->nChunksPerBlock;
+
+			dev->readChunkWithTagsFromNAND(dev,chunk,NULL,&tags);
+						      
+			if(tags.sequenceNumber == YAFFS_SEQUENCE_CHECKPOINT_DATA){
+				/* Right kind of block */
+				dev->checkpointNextBlock = tags.objectId;
+				dev->checkpointCurrentBlock = i;
+				dev->checkpointBlockList[dev->blocksInCheckpoint] = i;
+				dev->blocksInCheckpoint++;
+				T(YAFFS_TRACE_CHECKPOINT,(TSTR("found checkpt block %d"TENDSTR),i));
+				return;
+			}
+		}
+
+	T(YAFFS_TRACE_CHECKPOINT,(TSTR("found no more checkpt blocks"TENDSTR)));
+
+	dev->checkpointNextBlock = -1;
+	dev->checkpointCurrentBlock = -1;
+}
+
+
 int yaffs_CheckpointOpen(yaffs_Device *dev, int forWriting)
 {
-	int i;
 	
 	/* Got the functions we need? */
 	if (!dev->writeChunkWithTagsToNAND ||
 	    !dev->readChunkWithTagsFromNAND ||
-	    !dev->eraseBlockInNAND)
+	    !dev->eraseBlockInNAND ||
+	    !dev->markNANDBlockBad)
 		return 0;
-	
-	/* Got a valid checkpoint data region? */
-	if(!yaffs_CheckpointSpaceOk(dev))
+
+	if(forWriting && !yaffs_CheckpointSpaceOk(dev))
 		return 0;
 			
 	if(!dev->checkpointBuffer)
 		dev->checkpointBuffer = YMALLOC(dev->nBytesPerChunk);
 	if(!dev->checkpointBuffer)
 		return 0;
+
 	
-	dev->checkpointPage = dev->checkpointStartBlock * dev->nChunksPerBlock;
+	dev->checkpointPageSequence = 0;
 	
 	dev->checkpointOpenForWrite = forWriting;
 	
 	dev->checkpointByteCount = 0;
+	dev->checkpointCurrentBlock = -1;
+	dev->checkpointCurrentChunk = -1;
+	dev->checkpointNextBlock = dev->startBlock;
 	
 	/* Erase all the blocks in the checkpoint area */
 	if(forWriting){
-		T(YAFFS_TRACE_CHECKPOINT,(TSTR("erasing checkpt data"TENDSTR)));
-		for( i = dev->checkpointStartBlock; i <= dev->checkpointEndBlock; i++)
-			dev->eraseBlockInNAND(dev,i);
+		memset(dev->checkpointBuffer,0,dev->nBytesPerChunk);
+		dev->checkpointByteOffset = 0;
+		return yaffs_CheckpointErase(dev);
 		
-		dev->checkpointByteOffset = 0;
-		
-		memset(dev->checkpointBuffer,0,dev->nBytesPerChunk);
 		
 	} else {
+		int i;
 		/* Set to a value that will kick off a read */
 		dev->checkpointByteOffset = dev->nBytesPerChunk;
+		/* A checkpoint block list of 1 checkpoint block per 16 block is (hopefully)
+		 * going to be way more than we need */
+		dev->checkpointMaxBlocks = (dev->endBlock - dev->startBlock)/16 + 2;
+		dev->checkpointBlockList = YMALLOC(sizeof(int) * dev->checkpointMaxBlocks);
+		for(i = 0; i < dev->checkpointMaxBlocks; i++)
+			dev->checkpointBlockList[i] = -1;
 	}
 	
 	return 1;
@@ -86,47 +179,40 @@ int yaffs_CheckpointOpen(yaffs_Device *d
 static int yaffs_CheckpointFlushBuffer(yaffs_Device *dev)
 {
 
-	int blockNo;
-	int goodBlockFound;
-	yaffs_BlockState state;
-	__u32 seqenceNumber;
+	int chunk;
 
 	yaffs_ExtendedTags tags;
+	
+	if(dev->checkpointCurrentBlock < 0){
+		yaffs_CheckpointFindNextErasedBlock(dev);
+		dev->checkpointCurrentChunk = 0;
+	}
+	
+	if(dev->checkpointCurrentBlock < 0)
+		return 0;
+	
 	tags.chunkDeleted = 0;
-	tags.objectId = YAFFS_OBJECTID_CHECKPOINT_DATA;
-	tags.chunkId = dev->checkpointPage + 1;
-	tags.sequenceNumber = 1;
+	tags.objectId = dev->checkpointNextBlock; /* Hint to next place to look */
+	tags.chunkId = dev->checkpointPageSequence + 1;
+	tags.sequenceNumber =  YAFFS_SEQUENCE_CHECKPOINT_DATA;
 	tags.byteCount = dev->nBytesPerChunk;
 	
-	/* printf("write checkpoint page %d\n",dev->checkpointPage); */
-	
-	if(dev->checkpointPage%dev->nChunksPerBlock == 0){
-		/* Start of a new block, do a block validity check */
-		blockNo = dev->checkpointPage/dev->nChunksPerBlock;
-		goodBlockFound = 0;
-		while(blockNo <= dev->checkpointEndBlock && !goodBlockFound){
-			dev->queryNANDBlock(dev,blockNo,&state,&tags);
-			if(state != YAFFS_BLOCK_STATE_DEAD)
-				goodBlockFound = 1;
-			else {
-				blockNo++;
-				dev->checkpointPage += dev->nChunksPerBlock;
-			}
-		}
+	chunk = dev->checkpointCurrentBlock * dev->nChunksPerBlock + dev->checkpointCurrentChunk;
 		
-		if(!goodBlockFound)
-			return 0;
+	dev->writeChunkWithTagsToNAND(dev,chunk,dev->checkpointBuffer,&tags);
+	dev->checkpointByteOffset = 0;
+	dev->checkpointPageSequence++;	   
+	dev->checkpointCurrentChunk++;
+	if(dev->checkpointCurrentChunk >= dev->nChunksPerBlock){
+		dev->checkpointCurrentChunk = 0;
+		dev->checkpointCurrentBlock = -1;
 	}
-	
-	dev->writeChunkWithTagsToNAND(dev, dev->checkpointPage, 
-				      dev->checkpointBuffer,&tags);
-	dev->checkpointByteOffset = 0;
-	dev->checkpointPage++;	   
 	memset(dev->checkpointBuffer,0,dev->nBytesPerChunk);
 	
 	return 1;
 }
 
+
 int yaffs_CheckpointWrite(yaffs_Device *dev,const void *data, int nBytes)
 {
 	int i=0;
@@ -165,9 +251,9 @@ int yaffs_CheckpointRead(yaffs_Device *d
 	int i=0;
 	int ok = 1;
 	yaffs_ExtendedTags tags;
-	yaffs_BlockState state;
-	int blockNo;
-	int goodBlockFound;
+
+	
+	int chunk;
 
 	__u8 *dataBytes = (__u8 *)data;
 		
@@ -175,39 +261,40 @@ int yaffs_CheckpointRead(yaffs_Device *d
 		return 0;
 
 	while(i < nBytes && ok) {
+	
+	
 		if(dev->checkpointByteOffset < 0 ||
 		   dev->checkpointByteOffset >= dev->nBytesPerChunk) {
-			if(dev->checkpointPage%dev->nChunksPerBlock == 0){
-				/* Start of a new block, do a block validity check */
-				blockNo = dev->checkpointPage/dev->nChunksPerBlock;
-				goodBlockFound = 0;
-				while(blockNo <= dev->checkpointEndBlock && !goodBlockFound){
-					dev->queryNANDBlock(dev,blockNo,&state,&tags);
-					if(state != YAFFS_BLOCK_STATE_DEAD)
-						goodBlockFound = 1;
-					else {
-						blockNo++;
-						dev->checkpointPage += dev->nChunksPerBlock;
-					}
-				}
-		
-				if(!goodBlockFound)
-					return 0;
+		   
+		   	if(dev->checkpointCurrentBlock < 0){
+				yaffs_CheckpointFindNextCheckpointBlock(dev);
+				dev->checkpointCurrentChunk = 0;
 			}
-
+			
+			if(dev->checkpointCurrentBlock < 0)
+				ok = 0;
+			else {
+			
+				chunk = dev->checkpointCurrentBlock * dev->nChunksPerBlock + 
+				          dev->checkpointCurrentChunk;
 
-	   		/* read in the next chunk */
-	   		/* printf("read checkpoint page %d\n",dev->checkpointPage); */
-			dev->readChunkWithTagsFromNAND(dev, dev->checkpointPage, 
-						       dev->checkpointBuffer,
-						      &tags);
+	   			/* read in the next chunk */
+	   			/* printf("read checkpoint page %d\n",dev->checkpointPage); */
+				dev->readChunkWithTagsFromNAND(dev, chunk, 
+							       dev->checkpointBuffer,
+							      &tags);
 						      
-			if(tags.objectId != YAFFS_OBJECTID_CHECKPOINT_DATA ||
-			   tags.chunkId != (dev->checkpointPage + 1))
-			   ok = 0;
+				if(tags.chunkId != (dev->checkpointPageSequence + 1) ||
+				   tags.sequenceNumber != YAFFS_SEQUENCE_CHECKPOINT_DATA)
+				   ok = 0;
 
-			dev->checkpointByteOffset = 0;
-			dev->checkpointPage++;
+				dev->checkpointByteOffset = 0;
+				dev->checkpointPageSequence++;
+				dev->checkpointCurrentChunk++;
+			
+				if(dev->checkpointCurrentChunk >= dev->nChunksPerBlock)
+					dev->checkpointCurrentBlock = -1;
+			}
 		}
 		
 		if(ok){
@@ -228,13 +315,29 @@ int yaffs_CheckpointClose(yaffs_Device *
 	if(dev->checkpointOpenForWrite){	
 		if(dev->checkpointByteOffset != 0)
 			yaffs_CheckpointFlushBuffer(dev);
+	} else {
+		int i;
+		for(i = 0; i < dev->blocksInCheckpoint && dev->checkpointBlockList[i] >= 0; i++){
+			yaffs_BlockInfo *bi = &dev->blockInfo[dev->checkpointBlockList[i]];
+			if(bi->blockState == YAFFS_BLOCK_STATE_EMPTY)
+				bi->blockState = YAFFS_BLOCK_STATE_CHECKPOINT;
+			else {
+				// Todo this looks odd...
+			}
+		}
+		YFREE(dev->checkpointBlockList);
+		dev->checkpointBlockList = NULL;
 	}
 
+	dev->nFreeChunks -= dev->blocksInCheckpoint * dev->nChunksPerBlock;
+	dev->nErasedBlocks -= dev->blocksInCheckpoint;
+
+		
 	T(YAFFS_TRACE_CHECKPOINT,(TSTR("checkpoint byte count %d" TENDSTR),
 			dev->checkpointByteCount));
 			
 	if(dev->checkpointBuffer){
-		/* flush the buffer */	
+		/* free the buffer */	
 		YFREE(dev->checkpointBuffer);
 		dev->checkpointBuffer = NULL;
 		return 1;
@@ -253,13 +356,7 @@ int yaffs_CheckpointInvalidateStream(yaf
 	if(!yaffs_CheckpointSpaceOk(dev))
 		return 0;
 
-	if(dev->eraseBlockInNAND){
-		T(YAFFS_TRACE_CHECKPOINT,(TSTR("erasing checkpt data"TENDSTR)));
-
-		return dev->eraseBlockInNAND(dev, dev->checkpointStartBlock);
-	}
-	else
-		return 0;
+	return yaffs_CheckpointErase(dev);
 }