core/logic/smn_adt_trie.cpp
author Kyle Sanderson <kyle.leet@gmail.com>
Mon Aug 12 00:42:56 2013 +0100 (2013-08-12)
changeset 3994 4bc73bac0f4c
parent 2728 96022dad7e53
child 4096 1ed7fe4e8f7e
permissions -rw-r--r--
Added helper stock to check if a command exists. (bug 5510, r=asherkin)
     1 /**
     2  * vim: set ts=4 :
     3  * =============================================================================
     4  * SourceMod
     5  * Copyright (C) 2004-2008 AlliedModders LLC.  All rights reserved.
     6  * =============================================================================
     7  *
     8  * This program is free software; you can redistribute it and/or modify it under
     9  * the terms of the GNU General Public License, version 3.0, as published by the
    10  * Free Software Foundation.
    11  * 
    12  * This program is distributed in the hope that it will be useful, but WITHOUT
    13  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
    14  * FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
    15  * details.
    16  *
    17  * You should have received a copy of the GNU General Public License along with
    18  * this program.  If not, see <http://www.gnu.org/licenses/>.
    19  *
    20  * As a special exception, AlliedModders LLC gives you permission to link the
    21  * code of this program (as well as its derivative works) to "Half-Life 2," the
    22  * "Source Engine," the "SourcePawn JIT," and any Game MODs that run on software
    23  * by the Valve Corporation.  You must obey the GNU General Public License in
    24  * all respects for all other code used.  Additionally, AlliedModders LLC grants
    25  * this exception to all derivative works.  AlliedModders LLC defines further
    26  * exceptions, found in LICENSE.txt (as of this writing, version JULY-31-2007),
    27  * or <http://www.sourcemod.net/license.php>.
    28  *
    29  * Version: $Id$
    30  */
    31 
    32 #include <stdlib.h>
    33 #include "common_logic.h"
    34 #include <sm_trie_tpl.h>
    35 
    36 HandleType_t htCellTrie;
    37 
    38 enum TrieNodeType
    39 {
    40 	TrieNode_Cell,
    41 	TrieNode_CellArray,
    42 	TrieNode_String,
    43 };
    44 
    45 struct SmartTrieNode
    46 {
    47 	SmartTrieNode()
    48 	{
    49 		ptr = NULL;
    50 		type = TrieNode_Cell;
    51 	}
    52 	SmartTrieNode(const SmartTrieNode &obj)
    53 	{
    54 		type = obj.type;
    55 		ptr = obj.ptr;
    56 		data = obj.data;
    57 		data_len = obj.data_len;
    58 	}
    59 	SmartTrieNode & operator =(const SmartTrieNode &src)
    60 	{
    61 		type = src.type;
    62 		ptr = src.ptr;
    63 		data = src.data;
    64 		data_len = src.data_len;
    65 		return *this;
    66 	}
    67 	TrieNodeType type;
    68 	cell_t *ptr;
    69 	cell_t data;
    70 	cell_t data_len;
    71 };
    72 
    73 struct CellTrie
    74 {
    75 	KTrie<SmartTrieNode> trie;
    76 	cell_t mem_usage;
    77 };
    78 
    79 class TrieHelpers : 
    80 	public SMGlobalClass,
    81 	public IHandleTypeDispatch
    82 {
    83 public: //SMGlobalClass
    84 	void OnSourceModAllInitialized()
    85 	{
    86 		htCellTrie = handlesys->CreateType("Trie", this, 0, NULL, NULL, g_pCoreIdent, NULL);
    87 	}
    88 	void OnSourceModShutdown()
    89 	{
    90 		handlesys->RemoveType(htCellTrie, g_pCoreIdent);
    91 	}
    92 public: //IHandleTypeDispatch
    93 	static void DestroySmartTrieNode(SmartTrieNode *pNode)
    94 	{
    95 		free(pNode->ptr);
    96 	}
    97 	void OnHandleDestroy(HandleType_t type, void *object)
    98 	{
    99 		CellTrie *pTrie = (CellTrie *)object;
   100 
   101 		pTrie->trie.run_destructor(DestroySmartTrieNode);
   102 
   103 		delete pTrie;
   104 	}
   105 	bool GetHandleApproxSize(HandleType_t type, void *object, unsigned int *pSize)
   106 	{
   107 		CellTrie *pArray = (CellTrie *)object;
   108 		*pSize = sizeof(CellTrie) + pArray->mem_usage + pArray->trie.mem_usage();
   109 		return true;
   110 	}
   111 } s_CellTrieHelpers;
   112 
   113 static cell_t CreateTrie(IPluginContext *pContext, const cell_t *params)
   114 {
   115 	CellTrie *pTrie = new CellTrie;
   116 	Handle_t hndl;
   117 
   118 	pTrie->mem_usage = 0;
   119 
   120 	if ((hndl = handlesys->CreateHandle(htCellTrie, pTrie, pContext->GetIdentity(), g_pCoreIdent, NULL))
   121 		== BAD_HANDLE)
   122 	{
   123 		delete pTrie;
   124 		return BAD_HANDLE;
   125 	}
   126 
   127 	return hndl;
   128 }
   129 
   130 static void UpdateNodeCells(CellTrie *pTrie, SmartTrieNode *pData, const cell_t *cells, cell_t num_cells)
   131 {
   132 	if (num_cells == 1)
   133 	{
   134 		pData->data = *cells;
   135 		pData->type = TrieNode_Cell;
   136 	}
   137 	else
   138 	{
   139 		pData->type = TrieNode_CellArray;
   140 		if (pData->ptr == NULL)
   141 		{
   142 			pData->ptr = (cell_t *)malloc(num_cells * sizeof(cell_t));
   143 			pData->data_len = num_cells;
   144 			pTrie->mem_usage += (pData->data_len * sizeof(cell_t));
   145 		}
   146 		else if (pData->data_len < num_cells)
   147 		{
   148 			pData->ptr = (cell_t *)realloc(pData->ptr, num_cells * sizeof(cell_t));
   149 			pTrie->mem_usage += (num_cells - pData->data_len) * sizeof(cell_t);
   150 			pData->data_len = num_cells;
   151 		}
   152 		if (num_cells != 0)
   153 		{
   154 			memcpy(pData->ptr, cells, sizeof(cell_t) * num_cells);
   155 		}
   156 		pData->data = num_cells;
   157 	}
   158 }
   159 
   160 static void UpdateNodeString(CellTrie *pTrie, SmartTrieNode *pData, const char *str)
   161 {
   162 	size_t len = strlen(str);
   163 	cell_t num_cells = (len + sizeof(cell_t)) / sizeof(cell_t);
   164 
   165 	if (pData->ptr == NULL)
   166 	{
   167 		pData->ptr = (cell_t *)malloc(num_cells * sizeof(cell_t));
   168 		pData->data_len = num_cells;
   169 		pTrie->mem_usage += (pData->data_len * sizeof(cell_t));
   170 	}
   171 	else if (pData->data_len < num_cells)
   172 	{
   173 		pData->ptr = (cell_t *)realloc(pData->ptr, num_cells * sizeof(cell_t));
   174 		pTrie->mem_usage += (num_cells - pData->data_len) * sizeof(cell_t);
   175 		pData->data_len = num_cells;
   176 	}
   177 
   178 	strcpy((char *)pData->ptr, str);
   179 	pData->data = len;
   180 	pData->type = TrieNode_String;
   181 }
   182 
   183 static cell_t SetTrieValue(IPluginContext *pContext, const cell_t *params)
   184 {
   185 	Handle_t hndl;
   186 	CellTrie *pTrie;
   187 	HandleError err;
   188 	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);
   189 
   190 	hndl = params[1];
   191 
   192 	if ((err = handlesys->ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
   193 		!= HandleError_None)
   194 	{
   195 		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
   196 	}
   197 
   198 	char *key;
   199 	pContext->LocalToString(params[2], &key);
   200 
   201 	SmartTrieNode *pNode;
   202 	if ((pNode = pTrie->trie.retrieve(key)) == NULL)
   203 	{
   204 		SmartTrieNode node;
   205 		UpdateNodeCells(pTrie, &node, &params[3], 1);
   206 		return pTrie->trie.insert(key, node) ? 1 : 0;
   207 	}
   208 
   209 	if (!params[4])
   210 	{
   211 		return 0;
   212 	}
   213 
   214 	UpdateNodeCells(pTrie, pNode, &params[3], 1);
   215 
   216 	return 1;
   217 }
   218 
   219 static cell_t SetTrieArray(IPluginContext *pContext, const cell_t *params)
   220 {
   221 	Handle_t hndl;
   222 	CellTrie *pTrie;
   223 	HandleError err;
   224 	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);
   225 
   226 	hndl = params[1];
   227 
   228 	if ((err = handlesys->ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
   229 		!= HandleError_None)
   230 	{
   231 		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
   232 	}
   233 
   234 	if (params[4] < 0)
   235 	{
   236 		return pContext->ThrowNativeError("Invalid array size: %d", params[4]);
   237 	}
   238 
   239 	char *key;
   240 	cell_t *array;
   241 	pContext->LocalToString(params[2], &key);
   242 	pContext->LocalToPhysAddr(params[3], &array);
   243 
   244 	SmartTrieNode *pNode;
   245 	if ((pNode = pTrie->trie.retrieve(key)) == NULL)
   246 	{
   247 		SmartTrieNode node;
   248 		UpdateNodeCells(pTrie, &node, array, params[4]);
   249 		if (!pTrie->trie.insert(key, node))
   250 		{
   251 			free(node.ptr);
   252 			return 0;
   253 		}
   254 		return 1;
   255 	}
   256 
   257 	if (!params[4])
   258 	{
   259 		return 0;
   260 	}
   261 
   262 	UpdateNodeCells(pTrie, pNode, array, params[4]);
   263 
   264 	return 1;
   265 }
   266 
   267 static cell_t SetTrieString(IPluginContext *pContext, const cell_t *params)
   268 {
   269 	Handle_t hndl;
   270 	CellTrie *pTrie;
   271 	HandleError err;
   272 	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);
   273 
   274 	hndl = params[1];
   275 
   276 	if ((err = handlesys->ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
   277 		!= HandleError_None)
   278 	{
   279 		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
   280 	}
   281 
   282 	char *key, *val;
   283 	pContext->LocalToString(params[2], &key);
   284 	pContext->LocalToString(params[3], &val);
   285 
   286 	SmartTrieNode *pNode;
   287 	if ((pNode = pTrie->trie.retrieve(key)) == NULL)
   288 	{
   289 		SmartTrieNode node;
   290 		UpdateNodeString(pTrie, &node, val);
   291 		if (!pTrie->trie.insert(key, node))
   292 		{
   293 			free(node.ptr);
   294 			return 0;
   295 		}
   296 		return 1;
   297 	}
   298 
   299 	if (!params[4])
   300 	{
   301 		return 0;
   302 	}
   303 
   304 	UpdateNodeString(pTrie, pNode, val);
   305 
   306 	return 1;
   307 }
   308 
   309 static cell_t RemoveFromTrie(IPluginContext *pContext, const cell_t *params)
   310 {
   311 	Handle_t hndl;
   312 	CellTrie *pTrie;
   313 	HandleError err;
   314 	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);
   315 
   316 	hndl = params[1];
   317 
   318 	if ((err = handlesys->ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
   319 		!= HandleError_None)
   320 	{
   321 		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
   322 	}
   323 
   324 	char *key;
   325 	pContext->LocalToString(params[2], &key);
   326 
   327 	SmartTrieNode *pNode;
   328 	if ((pNode = pTrie->trie.retrieve(key)) == NULL)
   329 	{
   330 		return 0;
   331 	}
   332 
   333 	free(pNode->ptr);
   334 	pNode->ptr = NULL;
   335 
   336 	return pTrie->trie.remove(key) ? 1 : 0;
   337 }
   338 
   339 static cell_t ClearTrie(IPluginContext *pContext, const cell_t *params)
   340 {
   341 	Handle_t hndl;
   342 	CellTrie *pTrie;
   343 	HandleError err;
   344 	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);
   345 
   346 	hndl = params[1];
   347 
   348 	if ((err = handlesys->ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
   349 		!= HandleError_None)
   350 	{
   351 		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
   352 	}
   353 
   354 	pTrie->trie.run_destructor(TrieHelpers::DestroySmartTrieNode);
   355 	pTrie->trie.clear();
   356 
   357 	return 1;
   358 }
   359 
   360 static cell_t GetTrieValue(IPluginContext *pContext, const cell_t *params)
   361 {
   362 	Handle_t hndl;
   363 	CellTrie *pTrie;
   364 	HandleError err;
   365 	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);
   366 
   367 	hndl = params[1];
   368 
   369 	if ((err = handlesys->ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
   370 		!= HandleError_None)
   371 	{
   372 		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
   373 	}
   374 
   375 	char *key;
   376 	cell_t *pValue;
   377 	pContext->LocalToString(params[2], &key);
   378 	pContext->LocalToPhysAddr(params[3], &pValue);
   379 
   380 	SmartTrieNode *pNode;
   381 	if ((pNode = pTrie->trie.retrieve(key)) == NULL)
   382 	{
   383 		return 0;
   384 	}
   385 
   386 	if (pNode->type == TrieNode_Cell)
   387 	{
   388 		*pValue = pNode->data;
   389 		return 1;
   390 	}
   391 	
   392 	return 0;
   393 }
   394 
   395 static cell_t GetTrieArray(IPluginContext *pContext, const cell_t *params)
   396 {
   397 	Handle_t hndl;
   398 	CellTrie *pTrie;
   399 	HandleError err;
   400 	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);
   401 
   402 	hndl = params[1];
   403 
   404 	if ((err = handlesys->ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
   405 		!= HandleError_None)
   406 	{
   407 		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
   408 	}
   409 
   410 	if (params[4] < 0)
   411 	{
   412 		return pContext->ThrowNativeError("Invalid array size: %d", params[4]);
   413 	}
   414 
   415 	char *key;
   416 	cell_t *pValue, *pSize;
   417 	pContext->LocalToString(params[2], &key);
   418 	pContext->LocalToPhysAddr(params[3], &pValue);
   419 	pContext->LocalToPhysAddr(params[5], &pSize);
   420 
   421 	SmartTrieNode *pNode;
   422 	if ((pNode = pTrie->trie.retrieve(key)) == NULL
   423 		|| pNode->type != TrieNode_CellArray)
   424 	{
   425 		return 0;
   426 	}
   427 
   428 	if (pNode->ptr == NULL)
   429 	{
   430 		*pSize = 0;
   431 		return 1;
   432 	}
   433 
   434 	if (pNode->data > params[4])
   435 	{
   436 		*pSize = params[4];
   437 	}
   438 	else if (params[4] != 0)
   439 	{
   440 		*pSize = pNode->data;
   441 	}
   442 	else
   443 	{
   444 		return 1;
   445 	}
   446 
   447 	memcpy(pValue, pNode->ptr, sizeof(cell_t) * pSize[0]);
   448 
   449 	return 1;
   450 }
   451 
   452 static cell_t GetTrieString(IPluginContext *pContext, const cell_t *params)
   453 {
   454 	Handle_t hndl;
   455 	CellTrie *pTrie;
   456 	HandleError err;
   457 	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);
   458 
   459 	hndl = params[1];
   460 
   461 	if ((err = handlesys->ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
   462 		!= HandleError_None)
   463 	{
   464 		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
   465 	}
   466 
   467 	if (params[4] < 0)
   468 	{
   469 		return pContext->ThrowNativeError("Invalid buffer size: %d", params[4]);
   470 	}
   471 
   472 	char *key;
   473 	cell_t *pSize;
   474 	pContext->LocalToString(params[2], &key);
   475 	pContext->LocalToPhysAddr(params[5], &pSize);
   476 
   477 	SmartTrieNode *pNode;
   478 	if ((pNode = pTrie->trie.retrieve(key)) == NULL
   479 		|| pNode->type != TrieNode_String)
   480 	{
   481 		return 0;
   482 	}
   483 
   484 	if (pNode->ptr == NULL)
   485 	{
   486 		*pSize = 0;
   487 		pContext->StringToLocal(params[3], params[4], "");
   488 		return 1;
   489 	}
   490 
   491 	size_t written;
   492 	pContext->StringToLocalUTF8(params[3], params[4], (char *)pNode->ptr, &written);
   493 
   494 	*pSize = (cell_t)written;
   495 
   496 	return 1;
   497 }
   498 
   499 static cell_t GetTrieSize(IPluginContext *pContext, const cell_t *params)
   500 {
   501 	Handle_t hndl;
   502 	CellTrie *pTrie;
   503 	HandleError err;
   504 	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);
   505 
   506 	hndl = params[1];
   507 
   508 	if ((err = handlesys->ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
   509 		!= HandleError_None)
   510 	{
   511 		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
   512 	}
   513 
   514 	return pTrie->trie.size();
   515 }
   516 
   517 REGISTER_NATIVES(trieNatives)
   518 {
   519 	{"ClearTrie",				ClearTrie},
   520 	{"CreateTrie",				CreateTrie},
   521 	{"GetTrieArray",			GetTrieArray},
   522 	{"GetTrieString",			GetTrieString},
   523 	{"GetTrieValue",			GetTrieValue},
   524 	{"RemoveFromTrie",			RemoveFromTrie},
   525 	{"SetTrieArray",			SetTrieArray},
   526 	{"SetTrieString",			SetTrieString},
   527 	{"SetTrieValue",			SetTrieValue},
   528 	{"GetTrieSize",				GetTrieSize},
   529 	{NULL,						NULL},
   530 };