4 files changed, 160 insertions, 77 deletions

diff --git a/src/BlockTypePalette.cpp b/src/BlockTypePalette.cpp
index fabf5698e..f0f54b7c1 100644
--- a/src/BlockTypePalette.cpp
+++ b/src/BlockTypePalette.cpp
@@ -4,9 +4,9 @@
 
 
 
-BlockTypePalette::BlockTypePalette()
+BlockTypePalette::BlockTypePalette():
+	mMaxIndex(0)
 {
-	// Nothing needed yet
 }
 
 
@@ -22,8 +22,10 @@ UInt32 BlockTypePalette::index(const AString & aBlockTypeName, const BlockState
 	}
 
 	// Not found, append:
-	mPalette.push_back(std::make_pair(aBlockTypeName, aBlockState));
-	return static_cast<UInt32>(mPalette.size() - 1);
+	auto index = mMaxIndex++;
+	mBlockToNumber[aBlockTypeName][aBlockState] = index;
+	mNumberToBlock[index] = {aBlockTypeName, aBlockState};
+	return index;
 }
 
 
@@ -32,16 +34,17 @@ UInt32 BlockTypePalette::index(const AString & aBlockTypeName, const BlockState
 
 std::pair<UInt32, bool> BlockTypePalette::maybeIndex(const AString & aBlockTypeName, const BlockState & aBlockState) const
 {
-	auto count = mPalette.size();
-	for (size_t idx = 0; idx < count; ++idx)
+	auto itr1 = mBlockToNumber.find(aBlockTypeName);
+	if (itr1 == mBlockToNumber.end())
 	{
-		const auto & entry = mPalette[idx];
-		if ((entry.first == aBlockTypeName) && (entry.second == aBlockState))
-		{
-			return std::make_pair(static_cast<UInt32>(idx), true);
-		}
+		return {0, false};
+	}
+	auto itr2 = itr1->second.find(aBlockState);
+	if (itr2 == itr1->second.end())
+	{
+		return {0, false};
 	}
-	return std::make_pair(0, false);
+	return {itr2->second, true};
 }
 
 
@@ -50,7 +53,7 @@ std::pair<UInt32, bool> BlockTypePalette::maybeIndex(const AString & aBlockTypeN
 
 UInt32 BlockTypePalette::count() const
 {
-	return static_cast<UInt32>(mPalette.size());
+	return static_cast<UInt32>(mNumberToBlock.size());
 }
 
 
@@ -59,22 +62,54 @@ UInt32 BlockTypePalette::count() const
 
 const std::pair<AString, BlockState> & BlockTypePalette::entry(UInt32 aIndex) const
 {
-	ASSERT(aIndex < mPalette.size());
-	return mPalette[aIndex];
+	auto itr = mNumberToBlock.find(aIndex);
+	if (itr == mNumberToBlock.end())
+	{
+		throw NoSuchIndexException(aIndex);
+	}
+	return itr->second;
 }
 
 
 
 
 
-std::map<UInt32, UInt32> BlockTypePalette::createTransformMap(const BlockTypePalette & aOther)
+std::map<UInt32, UInt32> BlockTypePalette::createTransformMapAddMissing(const BlockTypePalette & aFrom)
 {
 	std::map<UInt32, UInt32> res;
-	auto numIndices = aOther.count();
-	for (UInt32 idx = 0; idx < numIndices; ++idx)
+	for (const auto & fromEntry: aFrom.mNumberToBlock)
 	{
-		const auto & e = aOther.mPalette[idx];
-		res[idx] = index(e.first, e.second);
+		auto fromIndex = fromEntry.first;
+		const auto & blockTypeName = fromEntry.second.first;
+		const auto & blockState = fromEntry.second.second;
+		res[fromIndex] = index(blockTypeName, blockState);
+	}
+	return res;
+}
+
+
+
+
+
+std::map<UInt32, UInt32> BlockTypePalette::createTransformMapWithFallback(const BlockTypePalette & aFrom, UInt32 aFallbackIndex) const
+{
+	std::map<UInt32, UInt32> res;
+	for (const auto & fromEntry: aFrom.mNumberToBlock)
+	{
+		auto fromIndex = fromEntry.first;
+		const auto & blockTypeName = fromEntry.second.first;
+		const auto & blockState = fromEntry.second.second;
+		auto thisIndex = maybeIndex(blockTypeName, blockState);
+		if (thisIndex.second)
+		{
+			// The entry was found in this
+			res[fromIndex] = thisIndex.first;
+		}
+		else
+		{
+			// The entry was NOT found in this, replace with fallback:
+			res[fromIndex] = aFallbackIndex;
+		}
 	}
 	return res;
 }
diff --git a/src/BlockTypePalette.h b/src/BlockTypePalette.h
index 47318f171..adb156ff1 100644
--- a/src/BlockTypePalette.h
+++ b/src/BlockTypePalette.h
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <unordered_map>
 #include <utility>
 #include "BlockState.h"
 
@@ -7,13 +8,33 @@
 
 
 
-/** Holds a palette that maps block type + state into numbers.
-Used primarily by PalettedBlockArea to translate between numeric and stringular block representation.
-The object itself provides no thread safety, users of this class need to handle locking, if required. */
+/** Holds a palette that maps between block type + state and numbers.
+Used primarily by PalettedBlockArea to map from stringular block representation to numeric,
+and by protocols to map from stringular block representation to protocol-numeric.
+The object itself provides no thread safety, users of this class need to handle locking, if required.
+Note that the palette itself doesn't support erasing;
+to erase, create a new instance and re-add only the wanted items.
+
+Internally, the object uses two synced maps, one for each translation direction. */
 class BlockTypePalette
 {
 public:
 
+	/** Exception that is thrown if requiesting an index not present in the palette. */
+	class NoSuchIndexException:
+		public std::runtime_error
+	{
+		using Super = std::runtime_error;
+
+	public:
+		NoSuchIndexException(UInt32 aIndex):
+			Super(Printf("No such palette index: %u", aIndex))
+		{
+		}
+	};
+
+
+
 	/** Create a new empty instance. */
 	BlockTypePalette();
 
@@ -29,16 +50,31 @@ public:
 	UInt32 count() const;
 
 	/** Returns the blockspec represented by the specified palette index.
-	The index must be valid (ASSERTed). */
+	If the index is not valid, throws a NoSuchIndexException. */
 	const std::pair<AString, BlockState> & entry(UInt32 aIndex) const;
 
-	/** Adds missing entries from aOther to this, and returns an index-transform map from aOther to this.
+	/** Returns an index-transform map from aFrom to this (this.entry(idx) == aFrom.entry(res[idx])).
+	Entries from aFrom that are not present in this are added.
 	Used when pasting two areas, to transform the src palette to dst palette. */
-	std::map<UInt32, UInt32> createTransformMap(const BlockTypePalette & aOther);
+	std::map<UInt32, UInt32> createTransformMapAddMissing(const BlockTypePalette & aFrom);
+
+	/** Returns an index-transform map from aFrom to this (this.entry(idx) == aFrom.entry(res[idx])).
+	Entries from aFrom that are not present in this are assigned the fallback index.
+	Used for protocol block type mapping. */
+	std::map<UInt32, UInt32> createTransformMapWithFallback(const BlockTypePalette & aFrom, UInt32 aFallbackIndex) const;
 
 
 protected:
 
-	/** The palette. Each item in the vector represents a single entry in the palette, the vector index is the palette index. */
-	std::vector<std::pair<AString, BlockState>> mPalette;
+	/** The mapping from numeric to stringular representation.
+	mNumberToBlock[index] = {"blockTypeName", blockState}. */
+	std::map<UInt32, std::pair<AString, BlockState>> mNumberToBlock;
+
+	/** The mapping from stringular to numeric representation.
+	mStringToNumber["blockTypeName"][blockState] = index. */
+	std::unordered_map<AString, std::map<BlockState, UInt32>> mBlockToNumber;
+
+	/** The maximum index ever used in the maps.
+	Used when adding new entries through the index() call. */
+	UInt32 mMaxIndex;
 };
diff --git a/src/PalettedBlockArea.cpp b/src/PalettedBlockArea.cpp
index e703adec0..fd1dcf332 100644
--- a/src/PalettedBlockArea.cpp
+++ b/src/PalettedBlockArea.cpp
@@ -178,7 +178,7 @@ void PalettedBlockArea::paste(const PalettedBlockArea & aSrc, const cCuboid & aS
 	}
 
 	// Create a transform map from aSrc's palette to our palette:
-	auto paletteTransform = mPalette.createTransformMap(aSrc.mPalette);
+	auto paletteTransform = mPalette.createTransformMapAddMissing(aSrc.mPalette);
 
 	// Copy the data:
 	UInt32 srcStrideY = static_cast<UInt32>(aSrc.size().x * aSrc.size().z);
diff --git a/tests/BlockTypeRegistry/BlockTypePaletteTest.cpp b/tests/BlockTypeRegistry/BlockTypePaletteTest.cpp
index ef79d8927..1337c7dc3 100644
--- a/tests/BlockTypeRegistry/BlockTypePaletteTest.cpp
+++ b/tests/BlockTypeRegistry/BlockTypePaletteTest.cpp
@@ -35,11 +35,12 @@ static void testBasic()
 	TEST_EQUAL(pal.count(), 5);
 
 	// Check the entry() API:
-	TEST_EQUAL(pal.entry(0), (std::make_pair<AString, BlockState>("testblock", BlockState())));
-	TEST_EQUAL(pal.entry(1), (std::make_pair<AString, BlockState>("another", BlockState())));
+	TEST_EQUAL(pal.entry(0), (std::make_pair<AString, BlockState>("testblock",  BlockState())));
+	TEST_EQUAL(pal.entry(1), (std::make_pair<AString, BlockState>("another",    BlockState())));
 	TEST_EQUAL(pal.entry(2), (std::make_pair<AString, BlockState>("multistate", BlockState(bs1))));  // make_pair requires a copy of the state
 	TEST_EQUAL(pal.entry(3), (std::make_pair<AString, BlockState>("multistate", BlockState(bs2))));
 	TEST_EQUAL(pal.entry(4), (std::make_pair<AString, BlockState>("multistate", BlockState(bs3))));
+	TEST_THROWS(pal.entry(5), BlockTypePalette::NoSuchIndexException);
 }
 
 
@@ -47,26 +48,26 @@ static void testBasic()
 
 
 /** Tests creating the transform map between two palettes. */
-static void testTransform()
+static void testTransformAddMissing()
 {
-	LOGD("Testing the createTransformMap API...");
+	LOGD("Testing the createTransformMapAddMissing API...");
 
 	// Create two palettes with some overlap:
 	BlockTypePalette pal1, pal2;
-	pal1.index("block1", BlockState());
-	pal1.index("block2", BlockState());
-	pal1.index("block3", BlockState());
-	pal1.index("block4", BlockState());
-	pal1.index("block5", BlockState("key1", "value1"));
-	pal2.index("block0", BlockState());
-	pal2.index("block2", BlockState());  // overlap
-	pal2.index("block3", BlockState());  // overlap
-	pal2.index("block4", BlockState("key1", "value1"));
-	pal2.index("block5", BlockState("key1", "value1"));  // overlap
-	pal2.index("block6", BlockState("key1", "value1"));
+	/* 0 */ pal1.index("block1", BlockState());
+	/* 1 */ pal1.index("block2", BlockState());
+	/* 2 */ pal1.index("block3", BlockState());
+	/* 3 */ pal1.index("block4", BlockState());
+	/* 4 */ pal1.index("block5", BlockState("key1", "value1"));
+	/* 0 */ pal2.index("block0", BlockState());
+	/* 1 */ pal2.index("block2", BlockState());  // overlap
+	/* 2 */ pal2.index("block4", BlockState());  // overlap
+	/* 3 */ pal2.index("block4", BlockState("key1", "value1"));
+	/* 4 */ pal2.index("block5", BlockState("key1", "value1"));  // overlap
+	/* 5 */ pal2.index("block6", BlockState("key1", "value1"));
 
 	// Check the transform map:
-	auto trans = pal1.createTransformMap(pal2);
+	auto trans = pal1.createTransformMapAddMissing(pal2);
 	TEST_EQUAL(pal1.maybeIndex("block1", BlockState()),                 (std::make_pair<UInt32, bool>(0, true)));
 	TEST_EQUAL(pal1.maybeIndex("block2", BlockState()),                 (std::make_pair<UInt32, bool>(1, true)));
 	TEST_EQUAL(pal1.maybeIndex("block3", BlockState()),                 (std::make_pair<UInt32, bool>(2, true)));
@@ -76,47 +77,58 @@ static void testTransform()
 	TEST_EQUAL(pal1.maybeIndex("block4", BlockState("key1", "value1")), (std::make_pair<UInt32, bool>(6, true)));
 	TEST_EQUAL(pal1.maybeIndex("block6", BlockState("key1", "value1")), (std::make_pair<UInt32, bool>(7, true)));
 	TEST_EQUAL(trans.size(), 6);
-	TEST_EQUAL(trans[0], 5);
-	TEST_EQUAL(trans[1], 1);
-	TEST_EQUAL(trans[2], 2);
-	TEST_EQUAL(trans[3], 6);
-	TEST_EQUAL(trans[4], 4);
-	TEST_EQUAL(trans[5], 7);
+	TEST_EQUAL(trans[0], 5);  // Added
+	TEST_EQUAL(trans[1], 1);  // Mapped
+	TEST_EQUAL(trans[2], 3);  // Mapped
+	TEST_EQUAL(trans[3], 6);  // Added
+	TEST_EQUAL(trans[4], 4);  // Mapped
+	TEST_EQUAL(trans[5], 7);  // Added
 }
 
 
 
 
 
-int main()
+/** Tests creating the transform map between two palettes, with fallback. */
+static void testTransformWithFallback()
 {
-	LOGD("BlockTypePaletteTest started");
-
-	try
-	{
-		testBasic();
-		testTransform();
-	}
-	catch (const TestException & exc)
-	{
-		LOGERROR("BlockTypePaletteTest has failed, an exception was thrown: %s", exc.mMessage.c_str());
-		return 1;
-	}
-	catch (const std::exception & exc)
-	{
-		LOGERROR("BlockTypePaletteTest has failed, an exception was thrown: %s", exc.what());
-		return 1;
-	}
-	catch (...)
-	{
-		LOGERROR("BlockTypePaletteTest has failed, an unhandled exception was thrown.");
-		return 1;
-	}
-
-	LOGD("BlockTypePaletteTest finished");
-	return 0;
+	LOGD("Testing the createTransformMapWithFallback API...");
+
+	// Create two palettes with some overlap:
+	BlockTypePalette pal1, pal2;
+	/* 0 */ pal1.index("block1", BlockState());
+	/* 1 */ pal1.index("block2", BlockState());
+	/* 2 */ pal1.index("block3", BlockState());
+	/* 3 */ pal1.index("block4", BlockState());
+	/* 4 */ pal1.index("block5", BlockState("key1", "value1"));
+	/* 0 */ pal2.index("block0", BlockState());
+	/* 1 */ pal2.index("block2", BlockState());  // overlap
+	/* 2 */ pal2.index("block4", BlockState());  // overlap
+	/* 3 */ pal2.index("block4", BlockState("key1", "value1"));
+	/* 4 */ pal2.index("block5", BlockState("key1", "value1"));  // overlap
+	/* 5 */ pal2.index("block6", BlockState("key1", "value1"));
+
+	// Check the transform map:
+	auto trans = pal1.createTransformMapWithFallback(pal2, 0);
+	TEST_EQUAL(trans.size(), 6);
+	TEST_EQUAL(trans[0], 0);  // Fallback
+	TEST_EQUAL(trans[1], 1);  // Mapped
+	TEST_EQUAL(trans[2], 3);  // Mapped
+	TEST_EQUAL(trans[3], 0);  // Fallback
+	TEST_EQUAL(trans[4], 4);  // Mapped
+	TEST_EQUAL(trans[5], 0);  // Fallback
 }
 
 
 
 
+
+IMPLEMENT_TEST_MAIN("BlockTypePalette",
+	testBasic();
+	testTransformAddMissing();
+	testTransformWithFallback();
+)
+
+
+
+