3 files changed, 100 insertions, 29 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);