scriptcomp: Support constant folding on unary operators too (#87)

Fixes #69. Negation was already handled for consts as a special case that we _might_ be able to remove now, but I'm leaving it for now. ## Testing Added a few tests. ## Licence - [x] I am licencing my change under the project's MIT licence, including all changes to GPL-3.0 licenced parts of the codebase.
niv · Dec 25, 2023 · a37952e · a37952e
1 parent ae7b436
commit a37952e
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 9 deletions.
diff --git a/neverwinter/nwscript/native/scriptcompcore.cpp b/neverwinter/nwscript/native/scriptcompcore.cpp
@@ -1527,15 +1527,19 @@ BOOL CScriptCompiler::ConstantFoldNode(CScriptParseTreeNode *pNode, BOOL bForce)
 	if (!pNode)
 		return FALSE;
 
-	// Only fold operations that have two operands
-	// TODO: ~0 unary op?
-	if (!pNode->pLeft || !pNode->pRight)
+	BOOL bUnary = pNode->nOperation == CSCRIPTCOMPILER_OPERATION_BOOLEAN_NOT ||
+	              pNode->nOperation == CSCRIPTCOMPILER_OPERATION_ONES_COMPLEMENT ||
+	              pNode->nOperation == CSCRIPTCOMPILER_OPERATION_NEGATION;
+
+	// Only fold operations that have all operands
+	if (!pNode->pLeft || (!bUnary && !pNode->pRight))
 		return FALSE;
 
 	// In case of complex expression, start folding at the leaf nodes
 	// e.g.:  C = 3 + 2*4 - First fold 2*4 into 8, then 3+8 into 11
 	ConstantFoldNode(pNode->pLeft, bForce);
-	ConstantFoldNode(pNode->pRight, bForce);
+	if (pNode->pRight)
+		ConstantFoldNode(pNode->pRight, bForce);
 
 	// Can only fold if the operands are constants.
 	if (pNode->pLeft->nOperation != CSCRIPTCOMPILER_OPERATION_CONSTANT_INTEGER &&
@@ -1546,14 +1550,14 @@ BOOL CScriptCompiler::ConstantFoldNode(CScriptParseTreeNode *pNode, BOOL bForce)
 	}
 
 	// Only fold operations on same type. Expressions like "3.0f + 1" are not folded
-	if (pNode->pLeft->nOperation != pNode->pRight->nOperation)
+	if (pNode->pRight && (pNode->pLeft->nOperation != pNode->pRight->nOperation))
 		return FALSE;
 
 	if (pNode->pLeft->nOperation == CSCRIPTCOMPILER_OPERATION_CONSTANT_INTEGER)
 	{
 		int32_t result;
 		int32_t left = pNode->pLeft->nIntegerData;
-		int32_t right = pNode->pRight->nIntegerData;
+		int32_t right = pNode->pRight ? pNode->pRight->nIntegerData : 0;
 		switch (pNode->nOperation)
 		{
 			case CSCRIPTCOMPILER_OPERATION_LOGICAL_OR:          result = left || right; break;
@@ -1574,10 +1578,16 @@ BOOL CScriptCompiler::ConstantFoldNode(CScriptParseTreeNode *pNode, BOOL bForce)
 			case CSCRIPTCOMPILER_OPERATION_MULTIPLY:            result = left *  right; break;
 			case CSCRIPTCOMPILER_OPERATION_DIVIDE:              result = left /  right; break;
 			case CSCRIPTCOMPILER_OPERATION_MODULUS:             result = left %  right; break;
+			// Unary ops
+			case CSCRIPTCOMPILER_OPERATION_BOOLEAN_NOT:         result = !left; break;
+			case CSCRIPTCOMPILER_OPERATION_ONES_COMPLEMENT:     result = ~left; break;
+			case CSCRIPTCOMPILER_OPERATION_NEGATION:            result = -left; break;
+
 			default: return FALSE;
 		}
 		pNode->pLeft->Clean();
-		pNode->pRight->Clean();
+		if (pNode->pRight)
+			pNode->pRight->Clean();
 		pNode->Clean();
 		pNode->nOperation = CSCRIPTCOMPILER_OPERATION_CONSTANT_INTEGER;
 		pNode->nIntegerData = result;
@@ -1589,7 +1599,7 @@ BOOL CScriptCompiler::ConstantFoldNode(CScriptParseTreeNode *pNode, BOOL bForce)
 		float result;
 		int resultBool = -1;
 		float left = pNode->pLeft->fFloatData;
-		float right = pNode->pRight->fFloatData;
+		float right = pNode->pRight ? pNode->pRight->fFloatData : 0.0f;
 		switch (pNode->nOperation)
 		{
 			case CSCRIPTCOMPILER_OPERATION_ADD:                 result = left +  right; break;
@@ -1603,10 +1613,14 @@ BOOL CScriptCompiler::ConstantFoldNode(CScriptParseTreeNode *pNode, BOOL bForce)
 			case CSCRIPTCOMPILER_OPERATION_CONDITION_GT:        resultBool = left >  right; break;
 			case CSCRIPTCOMPILER_OPERATION_CONDITION_LT:        resultBool = left <  right; break;
 			case CSCRIPTCOMPILER_OPERATION_CONDITION_LEQ:       resultBool = left <= right; break;
+			// Unary ops
+			case CSCRIPTCOMPILER_OPERATION_NEGATION:            result = -left; break;
+
 			default: return FALSE;
 		}
 		pNode->pLeft->Clean();
-		pNode->pRight->Clean();
+		if (pNode->pRight)
+			pNode->pRight->Clean();
 		pNode->Clean();
 		if (resultBool != -1)
 		{

diff --git a/tests/scriptcomp/corpus/constants.nss b/tests/scriptcomp/corpus/constants.nss
@@ -38,6 +38,11 @@ const int CONSTINT_SUBTRACT            = A -  B;
 const int CONSTINT_MULTIPLY            = A *  B;
 const int CONSTINT_DIVIDE              = A /  B;
 const int CONSTINT_MODULUS             = A %  B;
+const int CONSTINT_BOOLEAN_NOT         = !A;
+const int CONSTINT_ONES_COMPLEMENT     = ~A;
+const int CONSTINT_NEGATION            = -A;
+
+const int CONSTINT_COMPLEX_EXPRESSION  = A * (B/A) + !A + !B;
 
 const float X = 10.0f;
 const float Y = X + X;
@@ -52,6 +57,7 @@ const int   CONSTFLOAT_CONDITION_GEQ       = X >= Y;
 const int   CONSTFLOAT_CONDITION_GT        = X >  Y;
 const int   CONSTFLOAT_CONDITION_LT        = X <  Y;
 const int   CONSTFLOAT_CONDITION_LEQ       = X <= Y;
+const float CONSTFLOAT_NEGATION            = -X;
 
 const string S1 = "AAA";
 const string S2 = "BBB";
@@ -112,6 +118,11 @@ void main()
     Assert(CONSTINT_MULTIPLY == 200);
     Assert(CONSTINT_DIVIDE == 0);
     Assert(CONSTINT_MODULUS == 10);
+    Assert(CONSTINT_BOOLEAN_NOT == 0);
+    Assert(CONSTINT_ONES_COMPLEMENT == -11);
+    Assert(CONSTINT_NEGATION == -10);
+
+    Assert(CONSTINT_COMPLEX_EXPRESSION == 20);
 
     Assert(S3 == "AAA_BBB");
     Assert(CONSTSTR_CONDITION_EQUAL == FALSE);