f3probe: add option --min-memory

This option instructs f3probe to trade speed for less memory usage,
that is, f3probe minimizes use of memory.

Currently, this option only drops the use of the bitmap of
the safe device.
Nevertheless, this is not negligible memory.
For example, a 1TB drive whose block size is 512 Bytes requires
256MB of RAM for this bitmap:

  1TB / 512Byte/Block = 2^31Block
  2^31Block / 1Block/bit = 2^31bit
  2^31bit / 8bit/Byte = 2^28Byte = 256MB

To put it in context, 256MB of RAM was all that
Raspberry Pi Model A had.

f3probe.c

@@ -43,6 +43,8 @@ static struct argp_option options[] = {
 		"Run a unit test; it ignores all other debug options",	0},
 	{"destructive",		'n',	NULL,		0,
 		"Do not restore blocks of the device after probing it",	2},
+	{"min-memory",		'l',	NULL,		0,
+		"Trade speed for less use of memory",		0},
 	{"manual-reset",	'm',	NULL,		0,
 		"Ask user to manually reset the drive",	0},
 	{ 0 }
@@ -52,18 +54,21 @@ struct args {
 	char		*filename;
 
 	/* Debugging options. */
-	bool		unit_test;
 	bool		debug;
+	bool		unit_test;
 	bool		keep_file;
-	bool		save;
 
+	/* Behavior options. */
+	bool		save;
+	bool		min_mem;
+	bool		manual_reset;
+	/* 2 free bytes. */
+
+	/* Geometry. */
 	uint64_t	real_size_byte;
 	uint64_t	fake_size_byte;
 	int		wrap;
 	int		block_order;
-
-	bool		manual_reset;
-	/* 3 free bytes. */
 };
 
 static long long arg_to_long_long(const struct argp_state *state,
@@ -137,6 +142,10 @@ static error_t parse_opt(int key, char *arg, struct argp_state *state)
 		args->save = false;
 		break;
 
+	case 'l':
+		args->min_mem = true;
+		break;
+
 	case 'm':
 		args->manual_reset = true;
 		break;
@@ -192,11 +201,14 @@ static const struct unit_test_item ftype_to_params[] = {
 	/* Geometry of a real limbo drive. */
 	{1777645568ULL,	32505331712ULL,	35,	9},
 
-	/* Geometry of a real wraparound drive. */
+	/* Wraparound drive. */
 	{1ULL << 31,	1ULL << 34,	31,	9},
 
 	/* Chain drive. */
 	{1ULL << 31,	1ULL << 34,	32,	9},
+
+	/* Extreme case for memory usage (limbo drive). */
+	{1ULL << 20,	1ULL << 40,	40,	9},
 };
 
 #define UNIT_TEST_N_CASES \
@@ -297,7 +309,17 @@ static int test_device(struct args *args)
 	assert(dev);
 
 	if (args->save) {
-		dev = create_safe_device(dev, probe_device_max_blocks(dev));
+		dev = create_safe_device(dev, probe_device_max_blocks(dev),
+			args->min_mem);
+		if (!dev) {
+			if (!args->min_mem)
+				fprintf(stderr, "Out of memory, try `f3probe --min-memory %s'\n",
+					args->filename);
+			else
+				fprintf(stderr, "Out of memory, try `f3probe --destructive %s'\nPlease back your data up before using option --destructive.\nAlternatively, you could use a machine with more memory to run f3probe.\n",
+					args->filename);
+			exit(1);
+		}
 		assert(dev);
 	}
 
@@ -346,15 +368,16 @@ int main(int argc, char **argv)
 {
 	struct args args = {
 		/* Defaults. */
-		.unit_test	= false,
 		.debug		= false,
+		.unit_test	= false,
 		.keep_file	= false,
 		.save		= true,
+		.min_mem	= false,
+		.manual_reset	= false,
 		.real_size_byte	= 1ULL << 31,
 		.fake_size_byte	= 1ULL << 34,
 		.wrap		= 31,
 		.block_order	= 9,
-		.manual_reset	= false,
 	};
 
 	/* Read parameters. */

libprobe.c

@@ -570,33 +570,59 @@ static inline void *align_512(void *p)
 	return (void *)(   (ip + 511) & ~511   );
 }
 
-static int sdev_write_block(struct device *dev, const char *buf, int length,
+static int sdev_save_block(struct safe_device *sdev,
-	uint64_t offset)
+	int length, uint64_t offset)
 {
-	struct safe_device *sdev = dev_sdev(dev);
 	const int block_order = dev_get_block_order(sdev->shadow_dev);
 	lldiv_t idx = lldiv(offset >> block_order, SDEV_BITMAP_BITS_PER_WORD);
 	SDEV_BITMAP_WORD set_bit = (SDEV_BITMAP_WORD)1 << idx.rem;
+	char *block;
+	int rc;
 
 	/* The current implementation doesn't support variable lengths. */
 	assert(length == dev_get_block_size(sdev->shadow_dev));
 
 	/* Is this block already saved? */
-	if (!(sdev->sb_bitmap[idx.quot] & set_bit)) {
+	if (!sdev->sb_bitmap) {
-		/* No. Save this block. */
+		int i;
-		char *block = (char *)align_512(sdev->saved_blocks) +
+		/* Running without bitmap. */
-			(sdev->sb_n << block_order);
+		for (i = 0; i < sdev->sb_n; i++)
-		int rc;
+			if (sdev->sb_offsets[i] == offset) {
+				/* The block at @offset is already saved. */
+				return 0;
+			}
+	} else if (sdev->sb_bitmap[idx.quot] & set_bit) {
+		/* The block at @offset is already saved. */
+		return 0;
+	}
+
+	/* The block at @offset hasn't been saved before. Save this block. */
 	assert(sdev->sb_n < sdev->sb_max);
+	block = (char *)align_512(sdev->saved_blocks) +
+		(sdev->sb_n << block_order);
 	rc = sdev->shadow_dev->read_block(sdev->shadow_dev, block,
 		length, offset);
 	if (rc)
 		return rc;
+
+	/* Bookkeeping. */
+	if (sdev->sb_bitmap)
 		sdev->sb_bitmap[idx.quot] |= set_bit;
 	sdev->sb_offsets[sdev->sb_n] = offset;
 	sdev->sb_n++;
+	return 0;
 }
 
+static int sdev_write_block(struct device *dev, const char *buf, int length,
+	uint64_t offset)
+{
+	struct safe_device *sdev = dev_sdev(dev);
+	int rc;
+
+	rc = sdev_save_block(sdev, length, offset);
+	if (rc)
+		return rc;
+
 	return sdev->shadow_dev->write_block(sdev->shadow_dev, buf,
 		length, offset);
 }
@@ -638,12 +664,11 @@ static void sdev_free(struct device *dev)
 	free_device(sdev->shadow_dev);
 }
 
-struct device *create_safe_device(struct device *dev, int max_blocks)
+struct device *create_safe_device(struct device *dev, int max_blocks,
+	int min_memory)
 {
 	struct safe_device *sdev;
 	const int block_order = dev_get_block_order(dev);
-	lldiv_t idx = lldiv(dev_get_size_byte(dev) >> block_order,
-		SDEV_BITMAP_BITS_PER_WORD);
 	uint64_t length;
 
 	sdev = malloc(sizeof(*sdev));
@@ -659,11 +684,18 @@ struct device *create_safe_device(struct device *dev, int max_blocks)
 	if (!sdev->sb_offsets)
 		goto saved_blocks;
 
-	length = (idx.quot + (idx.rem ? 1 : 0)) * sizeof(SDEV_BITMAP_WORD);
+	if (!min_memory) {
+		lldiv_t idx = lldiv(dev_get_size_byte(dev) >> block_order,
+			SDEV_BITMAP_BITS_PER_WORD);
+		length = (idx.quot + (idx.rem ? 1 : 0)) *
+			sizeof(SDEV_BITMAP_WORD);
 		sdev->sb_bitmap = malloc(length);
 		if (!sdev->sb_bitmap)
 			goto offsets;
 		memset(sdev->sb_bitmap, 0, length);
+	} else {
+		sdev->sb_bitmap = NULL;
+	}
 
 	sdev->shadow_dev = dev;
 	sdev->sb_n = 0;

libprobe.h

@@ -38,7 +38,8 @@ struct device *create_file_device(const char *filename,
 
 struct device *create_block_device(const char *filename, int manual_reset);
 
-struct device *create_safe_device(struct device *dev, int max_blocks);
+struct device *create_safe_device(struct device *dev, int max_blocks,
+	int min_memory);
 
 void free_device(struct device *dev);