Add new 'replace' and 'size' paramaters in some natives.

Fix some typos
Update the test plugin

amxmodx/trie_natives.cpp

@@ -29,7 +29,8 @@ static cell AMX_NATIVE_CALL TrieClear(AMX *amx, cell *params)
 	t->map.clear();
 	return 1;
 }
-// native TrieSetCell(Trie:handle, const key[], any:value);
+
+// native TrieSetCell(Trie:handle, const key[], any:value, bool:replace = true);
 static cell AMX_NATIVE_CALL TrieSetCell(AMX *amx, cell *params)
 {
 	CellTrie *t = g_TrieHandles.lookup(params[1]);
@@ -50,13 +51,22 @@ static cell AMX_NATIVE_CALL TrieSetCell(AMX *amx, cell *params)
 		{
 			return 0;
 		}
+
+		i->value.setCell(params[3]);
+		return 1;
+	}
+
+	// Old plugin doesn't have 'replace' parameter.
+	if (*params / sizeof(cell) == 4 && !params[4])
+	{
+		return 0;
 	}
 
 	i->value.setCell(params[3]);
-	
 	return 1;
 }
-// native TrieSetString(Trie:handle, const key[], const data[]);
+
+// native TrieSetString(Trie:handle, const key[], const data[], bool:replace = true);
 static cell AMX_NATIVE_CALL TrieSetString(AMX *amx, cell *params)
 {
 	CellTrie *t = g_TrieHandles.lookup(params[1]);
@@ -78,13 +88,22 @@ static cell AMX_NATIVE_CALL TrieSetString(AMX *amx, cell *params)
 		{
 			return 0;
 		}
+
+		i->value.setString(value);
+		return 1;
+	}
+
+	// Old plugin doesn't have 'replace' parameter.
+	if (*params / sizeof(cell) == 4 && !params[4])
+	{
+		return 0;
 	}
 
 	i->value.setString(value);
-
 	return 1;
 }
-// native TrieSetArray(Trie:handle, const key[], const any:buffer[], buffsize)
+
+// native TrieSetArray(Trie:handle, const key[], const any:buffer[], buffsize, bool:replace = true)
 static cell AMX_NATIVE_CALL TrieSetArray(AMX *amx, cell *params)
 {
 	CellTrie *t = g_TrieHandles.lookup(params[1]);
@@ -95,8 +114,15 @@ static cell AMX_NATIVE_CALL TrieSetArray(AMX *amx, cell *params)
 		return 0;
 	}
 
+	if (params[4] < 0)
+	{
+		LogError(amx, AMX_ERR_NATIVE, "Invalid array size (%d)", params[4]);
+		return 0;
+	}
+
 	int len;
 	const char *key = get_amxstring(amx, params[2], 0, len);
+	cell *ptr = get_amxaddr(amx, params[3]);
 
 	StringHashMap<Entry>::Insert i = t->map.findForAdd(key);
 	if (!i.found())
@@ -105,12 +131,22 @@ static cell AMX_NATIVE_CALL TrieSetArray(AMX *amx, cell *params)
 		{
 			return 0;
 		}
+
+		i->key = key;
+		i->value.setArray(ptr, params[4]);
+		return 1;
 	}
 
-	i->value.setArray(get_amxaddr(amx, params[3]), params[4]);
+	// Old plugin doesn't have 'replace' parameter.
+	if (*params / sizeof(cell) == 4 && !params[5])
+	{
+		return 0;
 
+	}
+	i->value.setArray(ptr, params[4]);
 	return 1;
 }
+
 // native bool:TrieGetCell(Trie:handle, const key[], &any:value);
 static cell AMX_NATIVE_CALL TrieGetCell(AMX *amx, cell *params)
 {
@@ -140,9 +176,10 @@ static cell AMX_NATIVE_CALL TrieGetCell(AMX *amx, cell *params)
 		return 1;
 	}
 
-	return 1;
+	return 0;
 }
-// native bool:TrieGetString(Trie:handle, const key[], buff[], len);
+
+// native bool:TrieGetString(Trie:handle, const key[], buff[], len, &size = 0);
 static cell AMX_NATIVE_CALL TrieGetString(AMX *amx, cell *params)
 {
 	CellTrie *t = g_TrieHandles.lookup(params[1]);
@@ -153,8 +190,15 @@ static cell AMX_NATIVE_CALL TrieGetString(AMX *amx, cell *params)
 		return 0;
 	}
 
+	if (params[4] < 0)
+	{
+		LogError(amx, AMX_ERR_NATIVE, "Invalid buffer size (%d)", params[4]);
+		return 0;
+	}
+
 	int len;
 	const char *key = get_amxstring(amx, params[2], 0, len);
+	cell *pSize = get_amxaddr(amx, params[5]);
 
 	StringHashMap<Entry>::Result r = t->map.find(key);
 	if (!r.found() || !r->value.isString())
@@ -162,11 +206,12 @@ static cell AMX_NATIVE_CALL TrieGetString(AMX *amx, cell *params)
 		return 0;
 	}
 
-	set_amxstring_utf8(amx, params[3], r->value.chars(), r->value.arrayLength(), params[4] + 1); // + EOS
+	*pSize = (cell)set_amxstring_utf8(amx, params[3], r->value.chars(), strlen(r->value.chars()), params[4] + 1); // + EOS
 
 	return 1;
 }
-// native bool:TrieGetArray(Trie:handle, const key[], any:buff[], len);
+
+// native bool:TrieGetArray(Trie:handle, const key[], any:buff[], len, &size = 0);
 static cell AMX_NATIVE_CALL TrieGetArray(AMX *amx, cell *params)
 {
 	CellTrie *t = g_TrieHandles.lookup(params[1]);
@@ -177,10 +222,16 @@ static cell AMX_NATIVE_CALL TrieGetArray(AMX *amx, cell *params)
 		return 0;
 	}
 
+	if (params[4] < 0)
+	{
+		LogError(amx, AMX_ERR_NATIVE, "Invalid array size (%d)", params[4]);
+		return 0;
+	}
+
 	int len;
 	const char *key = get_amxstring(amx, params[2], 0, len);
-	cell *ptr = get_amxaddr(amx, params[3]);
+	cell *pValue = get_amxaddr(amx, params[3]);
-	size_t size = params[4];
+	cell *pSize = get_amxaddr(amx, params[5]);
 
 	StringHashMap<Entry>::Result r = t->map.find(key);
 	if (!r.found() || !r->value.isArray())
@@ -188,17 +239,29 @@ static cell AMX_NATIVE_CALL TrieGetArray(AMX *amx, cell *params)
 		return 0;
 	}
 
+	if (!r->value.array())
+	{
+		*pSize = 0;
+		return 1;
+	}
+
+	if (!params[4])
+	{
+		return 0;
+	}
+
 	size_t length = r->value.arrayLength();
 	cell *base = r->value.array();
 
-	if (length <= size)
+	if (length > size_t(params[4]))
-	{
+		*pSize = params[4];
-		size = length;
+	else
-	}
+		*pSize = length;
 
-	memcpy(ptr, base, sizeof(cell)* size);
+	memcpy(pValue, base, sizeof(cell) * pSize[0]);
 	return 1;
 }
+
 // native bool:TrieKeyExists(Trie:handle, const key[]);
 static cell AMX_NATIVE_CALL TrieKeyExists(AMX *amx, cell *params)
 {

amxmodx/trie_natives.h

@@ -61,19 +61,19 @@ public:
 	}
 
 	size_t arrayLength() const {
-		//assert(isArray());
+		assert(isArray());
 		return raw()->length;
 	}
 	cell *array() const {
-		//assert(isArray());
+		assert(isArray());
 		return reinterpret_cast<cell *>(raw()->base());
 	}
 	char *chars() const {
-		//assert(isString());
+		assert(isString());
 		return reinterpret_cast<char *>(raw()->base());
 	}
 	cell cell_() const {
-		//assert(isCell());
+		assert(isCell());
 		return data_;
 	}
 
@@ -110,13 +110,13 @@ private:
 	}
 	void setType(EntryType aType) {
 		control_ = uintptr_t(raw()) | uintptr_t(aType);
-		//assert(type() == aType);
+		assert(type() == aType);
 	}
 	void setTypeAndPointer(EntryType aType, ArrayInfo *ptr) {
 		// malloc() should guarantee 8-byte alignment at worst
-		//assert((uintptr_t(ptr) & 0x3) == 0);
+		assert((uintptr_t(ptr) & 0x3) == 0);
 		control_ = uintptr_t(ptr) | uintptr_t(aType);
-		//assert(type() == aType);
+		assert(type() == aType);
 	}
 	EntryType type() const {
 		return (EntryType)(control_ & 0x3);

plugins/include/celltrie.inc

@@ -39,11 +39,12 @@ native TrieClear(Trie:handle);
  * @param handle	Map Handle.
  * @param key		Key string.
  * @param value		Value to store at this key.
+ * @param replace	If false, operation will fail if the key is already set.
  *
  * @return			True on success, false on failure.
  * @error			Invalid Handle.
  */
-native TrieSetCell(Trie:handle, const key[], any:value);
+native TrieSetCell(Trie:handle, const key[], any:value, bool:replace = true);
 
 /**
  * Sets a string value in a Map, either inserting a new entry or replacing an old one.
@@ -51,11 +52,12 @@ native TrieSetCell(Trie:handle, const key[], any:value);
  * @param handle	Map Handle.
  * @param key		Key string.
  * @param value		String to store.
+ * @param replace	If false, operation will fail if the key is already set.
  *
  * @return			True on success, false on failure.
  * @error			Invalid Handle.
  */
-native TrieSetString(Trie:handle, const key[], const value[]);
+native TrieSetString(Trie:handle, const key[], const value[], bool:replace = true);
 
 /**
  * Sets an array value in a Map, either inserting a new entry or replacing an old one.
@@ -64,11 +66,12 @@ native TrieSetString(Trie:handle, const key[], const value[]);
  * @param key		Key string.
  * @param buffer	Array to store.
  * @param size		Number of items in the array.
+ * @param replace	If false, operation will fail if the key is already set.
  *
  * @return			True on success, false on failure.
  * @error			Invalid Handle.
  */
-native TrieSetArray(Trie:handle, const key[], const any:buffer[], size);
+native TrieSetArray(Trie:handle, const key[], const any:buffer[], size, bool:replace = true);
 
 /**
  * Retrieves a value in a Map.
@@ -89,12 +92,13 @@ native bool:TrieGetCell(Trie:handle, const key[], &any:value);
  * @param key			Key string.
  * @param output		Buffer to store value.
  * @param outputsize	Maximum size of string buffer.
+ * @param size			Optional parameter to store the number of bytes written to the buffer.
  *
  * @return				True on success.  False if the key is not set, or the key is set 
  *						as a value or array (not a string).
  * @error				Invalid Handle.
  */
-native bool:TrieGetString(Trie:handle, const key[], output[], outputsize);
+native bool:TrieGetString(Trie:handle, const key[], output[], outputsize, &size = 0);
 
 /**
  * Retrieves an array in a Map.
@@ -103,12 +107,13 @@ native bool:TrieGetString(Trie:handle, const key[], output[], outputsize);
  * @param key			Key string.
  * @param output		Buffer to store array.
  * @param outputsize	Maximum size of array buffer.
+ * @param size			Optional parameter to store the number of elements written to the buffer.
  *
  * @return				True on success.  False if the key is not set, or the key is set 
  *						as a value or string (not an array).
  * @error				Invalid Handle.
  */
-native bool:TrieGetArray(Trie:handle, const key[], any:output[], outputsize);
+native bool:TrieGetArray(Trie:handle, const key[], any:output[], outputsize, &size = 0);
 
 /**
  * Removes a key entry from a Map.

plugins/testsuite/trietest.sma

@@ -62,30 +62,57 @@ public trietest()
 		}
 	}
 	
+	// Setting K42K without replace should fail.
+	new value;
+	if (TrieSetCell(t, "K42K", 999, false))				{ ok = false; server_print("set trie K42K should fail"); }
+	if (!TrieGetCell(t, "K42K", value) || value != 42)	{ ok = false; server_print("value at K42K not correct"); }
+	if (!TrieSetCell(t, "K42K", 999))					{ ok = false; server_print("set trie K42K = 999 should succeed"); }
+	if (!TrieGetCell(t, "K42K", value) || value != 999)	{ ok = false; server_print("value at K42K not correct"); }
+
 	// Check size is 100.
 	if (TrieGetSize(t) != 100)
 	{
-		ok = false;
+		server_print("Map size mismatch (1), expected: %d got: %d", 100, TrieGetSize(t)); ok = false;
-		server_print("Map size mismatch, expected: %d got: %d", 100, TrieGetSize(t));
 	}
 	
+	if (TrieGetCell(t, "cat", value))
+	{
+		server_print("trie should not have a cat."); ok = false;
+	}
+	
+	// Check that "K42K" is not a string or array.
+	new array[32];
+	new string[32];
+	if (TrieGetArray(t, "K42K", array, sizeof(array)) ||
+		TrieGetString(t, "K42K", string, charsmax(string)))
+	{
+		server_print("entry K42K should not be an array or string"); ok = false;
+	}
+ 
+	TrieClear(t);
+ 
+	if (TrieGetSize(t) != 0)
+	{
+		server_print("Map size mismatch (2), expected: %d got: %d", 0, TrieGetSize(t)); ok = false;
+	}
+	
+	TrieDestroy(t);
+	
 	if (ok)
 		pass("Cell tests");
 	else
 		fail("Cell tests");
 
-	TrieClear(t);
-	TrieDestroy(t);
 
 	t = TrieCreate();
 
 	if (t == oldhandle)
 		pass("Recycle handles");
-
 	else
 		fail("Recycle handles");
 
 	ok = true;
+	
 	for (new i = 0; i < 100; i++)
 	{
 		static val[32];
@@ -100,29 +127,107 @@ public trietest()
 		static val[32];
 		static exp[32];
 		formatex(exp, charsmax(exp), "V%dV", i);
-		if (!TrieGetString(t, key, val, charsmax(val)))
+		new size;
+		
+		if (!TrieGetString(t, key, val, charsmax(val), size))
 		{
 			server_print("TrieGetString(%d, '%s', %s) failed", t, key, val);
 			ok = false;
 		}
-
 		else if (!equal(val, exp))
 		{
 			server_print("val mismatch, key: '%s' expected: '%s' got: '%s'", key, exp, val);
 			ok = false;
 		}
-
+		if (size != strlen(exp))
+		{
+			server_print("returned size mismatch, key: '%s' expected: '%s' got: '%s'", key, strlen(exp), size);
+			ok = false;
+		}
 	}
+	
+	if (TrieGetCell(t, "K42K", value) ||
+		TrieGetArray(t, "K42K", array, sizeof(array)))
+	{
+		server_print("entry K42K should not be an array or string"); ok = false;
+	}
+  
+	if (TrieGetString(t, "cat", string, charsmax(string)))
+	{
+		server_print("trie should not have a cat."); ok = false;
+	}
+	
 	if (ok)
 		pass("String tests");
-
 	else
 		fail("String tests");
 
+		
+	ok = true;
+
+	new data[5] = { 93, 1, 2, 3, 4 };
+	
+	if (!TrieSetArray(t, "K42K", data, sizeof data))
+	{
+		server_print("K42K should be a string."); ok = false;
+	}
+	if (!TrieGetArray(t, "K42K", array, sizeof(array)))
+	{
+		server_print("K42K should be V42V."); ok = false;
+	}
+	for (new i = 0; i < sizeof data; i++) 
+	{
+		if (data[i] != array[i])
+		{
+			server_print("K42K slot %d should be %d, got %d", i, data[i], array[i]); ok = false;
+		}
+	}
+	if (TrieGetCell(t, "K42K", value) ||
+		TrieGetString(t, "K42K", string, charsmax(string)))
+	{
+		server_print("entry K42K should not be an array or string"); ok = false;
+	}	
+	if (!TrieSetArray(t, "K42K", data, 1))
+    {
+		server_print("couldn't set K42K to 1-entry array"); ok = false;
+	}
+	if (!TrieGetArray(t, "K42K", array, sizeof(array), value))
+	{
+		server_print("couldn't fetch 1-entry array"); ok = false;
+	}
+	if (value != 1)
+	{
+		server_print("array size mismatch (%d, expected %d)", value, 1); ok = false;
+	}
+	
+	if (ok)
+		pass("Array tests");
+	else
+		fail("Array tests");
+		
+		
+	ok = true;
+	
+	// Remove "K42K".
+	if (!TrieDeleteKey(t, "K42K"))
+	{
+		server_print("K42K should have been removed"); ok = false;
+	}
+	if (TrieDeleteKey(t, "K42K"))
+	{
+		server_print("K42K should not exist"); ok =false;
+	}
+	if (TrieGetCell(t, "K42K", value) ||
+		TrieGetArray(t, "K42K", array, sizeof(array)) ||
+		TrieGetString(t, "K42K", string, charsmax(string)))
+	{
+		server_print("map should not have a K42K"); ok = false;
+	}
+	
 	TrieDestroy(t);
 	
 	t = TrieCreate();
-	ok = true;
+	
 	for (new i = 0; i < 1000; i++)
 	{
 		formatex(key, charsmax(key), "!%d!", i);
@@ -134,21 +239,16 @@ public trietest()
 
 		if (!TrieKeyExists(t, key))
 		{
-			ok = false;
+			server_print("Key '%s' does not exist", key); ok = false;
-			server_print("Key '%s' does not exist", key);
 		}
-		else
+		else if (!TrieDeleteKey(t, key))
 		{
-			if (!TrieDeleteKey(t, key))
+			server_print("Key '%s' could not be deleted", key); ok = false;
-			{
-				server_print("Key '%s' could not be deleted", key);
-				ok = false;
-			}
 		}
 	}
+	
 	if (ok)
 		pass("Exists/Delete");
-
 	else
 		fail("Exists/Delete");