Fix some string operations (#895)

ark120202 · web-flow · commit 0d7c5cd7bf46 · 2020-07-13T20:04:29.000+02:00
* Fix string.slice, string.substring, string.substr

* Fix string.charAt and string.charCodeAt

* Fix string.indexOf with negative indexes

* Fix out-of-bounds string access being "" instead of undefined

* Fix array.join flacky test

* Feedback

* Replace createExpressionPlusOne with more general addToNumericExpression
diff --git a/src/LuaLib.ts b/src/LuaLib.ts
@@ -54,13 +54,19 @@ export enum LuaLibFeature {
     WeakSet = "WeakSet",
     SourceMapTraceBack = "SourceMapTraceBack",
     Spread = "Spread",
+    StringAccess = "StringAccess",
+    StringCharAt = "StringCharAt",
+    StringCharCodeAt = "StringCharCodeAt",
     StringConcat = "StringConcat",
     StringEndsWith = "StringEndsWith",
     StringPadEnd = "StringPadEnd",
     StringPadStart = "StringPadStart",
     StringReplace = "StringReplace",
+    StringSlice = "StringSlice",
     StringSplit = "StringSplit",
     StringStartsWith = "StringStartsWith",
+    StringSubstr = "StringSubstr",
+    StringSubstring = "StringSubstring",
     StringTrim = "StringTrim",
     StringTrimEnd = "StringTrimEnd",
     StringTrimStart = "StringTrimStart",
diff --git a/src/lualib/StringAccess.ts b/src/lualib/StringAccess.ts
@@ -0,0 +1,5 @@
+function __TS__StringAccess(this: string, index: number) {
+    if (index >= 0 && index < this.length) {
+        return string.sub(this, index + 1, index + 1);
+    }
+}
diff --git a/src/lualib/StringCharAt.ts b/src/lualib/StringCharAt.ts
@@ -0,0 +1,5 @@
+function __TS__StringCharAt(this: string, pos: number): string {
+    if (pos !== pos) pos = 0;
+    if (pos < 0) return "";
+    return string.sub(this, pos + 1, pos + 1);
+}
diff --git a/src/lualib/StringCharCodeAt.ts b/src/lualib/StringCharCodeAt.ts
@@ -0,0 +1,5 @@
+function __TS__StringCharCodeAt(this: string, index: number): number {
+    if (index !== index) index = 0;
+    if (index < 0) return NaN;
+    return string.byte(this, index + 1) ?? NaN;
+}
diff --git a/src/lualib/StringPadEnd.ts b/src/lualib/StringPadEnd.ts
@@ -13,5 +13,5 @@ function __TS__StringPadEnd(this: string, maxLength: number, fillString = " "):
         fillString += fillString.repeat(maxLength / fillString.length);
     }
 
-    return this + fillString.slice(0, Math.floor(maxLength));
+    return this + string.sub(fillString, 1, Math.floor(maxLength));
 }
diff --git a/src/lualib/StringPadStart.ts b/src/lualib/StringPadStart.ts
@@ -13,5 +13,5 @@ function __TS__StringPadStart(this: string, maxLength: number, fillString = " ")
         fillString += fillString.repeat(maxLength / fillString.length);
     }
 
-    return fillString.slice(0, Math.floor(maxLength)) + this;
+    return string.sub(fillString, 1, Math.floor(maxLength)) + this;
 }
diff --git a/src/lualib/StringSlice.ts b/src/lualib/StringSlice.ts
@@ -0,0 +1,9 @@
+function __TS__StringSlice(this: string, start?: number, end?: number): string {
+    if (start === undefined || start !== start) start = 0;
+    if (end !== end) end = 0;
+
+    if (start >= 0) start += 1;
+    if (end !== undefined && end < 0) end -= 1;
+
+    return string.sub(this, start, end);
+}
diff --git a/src/lualib/StringSubstr.ts b/src/lualib/StringSubstr.ts
@@ -0,0 +1,12 @@
+function __TS__StringSubstr(this: string, from: number, length?: number): string {
+    if (from !== from) from = 0;
+
+    if (length !== undefined) {
+        if (length !== length || length <= 0) return "";
+        length += from;
+    }
+
+    if (from >= 0) from += 1;
+
+    return string.sub(this, from, length);
+}
diff --git a/src/lualib/StringSubstring.ts b/src/lualib/StringSubstring.ts
@@ -0,0 +1,17 @@
+function __TS__StringSubstring(this: string, start: number, end?: number): string {
+    if (end !== end) end = 0;
+
+    if (end !== undefined && start > end) {
+        [start, end] = [end, start];
+    }
+
+    if (start >= 0) {
+        start += 1;
+    } else {
+        start = 1;
+    }
+
+    if (end !== undefined && end < 0) end = 0;
+
+    return string.sub(this, start, end);
+}
diff --git a/src/lualib/declarations/string.d.ts b/src/lualib/declarations/string.d.ts
@@ -1,5 +1,9 @@
 /** @noSelf */
 declare namespace string {
+    function byte(s: string, i?: number): number | undefined;
+    /** @tupleReturn */
+    function byte(s: string, i?: number, j?: number): number[];
+
     /** @tupleReturn */
     function gsub(
         source: string,
diff --git a/src/transformation/builtins/index.ts b/src/transformation/builtins/index.ts
@@ -2,6 +2,7 @@ import * as ts from "typescript";
 import * as lua from "../../LuaAST";
 import { assume } from "../../utils";
 import { TransformationContext } from "../context";
+import { createNaN } from "../utils/lua-ast";
 import { importLuaLibFeature, LuaLibFeature } from "../utils/lualib";
 import { getIdentifierSymbolId } from "../utils/symbols";
 import {
@@ -121,12 +122,7 @@ export function transformBuiltinIdentifierExpression(
 ): lua.Expression | undefined {
     switch (node.text) {
         case "NaN":
-            return lua.createBinaryExpression(
-                lua.createNumericLiteral(0),
-                lua.createNumericLiteral(0),
-                lua.SyntaxKind.DivisionOperator,
-                node
-            );
+            return createNaN(node);
 
         case "Infinity":
             const math = lua.createIdentifier("math");
diff --git a/src/transformation/builtins/string.ts b/src/transformation/builtins/string.ts
@@ -2,7 +2,7 @@ import * as ts from "typescript";
 import * as lua from "../../LuaAST";
 import { TransformationContext } from "../context";
 import { unsupportedProperty } from "../utils/diagnostics";
-import { createExpressionPlusOne } from "../utils/lua-ast";
+import { addToNumericExpression, createNaN, getNumberLiteralValue } from "../utils/lua-ast";
 import { LuaLibFeature, transformLuaLibFunction } from "../utils/lualib";
 import { PropertyCallExpression, transformArguments } from "../visitors/call";
 
@@ -30,13 +30,20 @@ export function transformStringPrototypeCall(
             return transformLuaLibFunction(context, LuaLibFeature.StringReplace, node, caller, ...params);
         case "concat":
             return transformLuaLibFunction(context, LuaLibFeature.StringConcat, node, caller, ...params);
-        case "indexOf":
+
+        case "indexOf": {
             const stringExpression = createStringCall(
                 "find",
                 node,
                 caller,
                 params[0],
-                params[1] ? createExpressionPlusOne(params[1]) : lua.createNilLiteral(),
+                params[1]
+                    ? // string.find handles negative indexes by making it relative to string end, but for indexOf it's the same as 0
+                      lua.createCallExpression(
+                          lua.createTableIndexExpression(lua.createIdentifier("math"), lua.createStringLiteral("max")),
+                          [addToNumericExpression(params[1], 1), lua.createNumericLiteral(1)]
+                      )
+                    : lua.createNilLiteral(),
                 lua.createBooleanLiteral(true)
             );
 
@@ -46,35 +53,38 @@ export function transformStringPrototypeCall(
                 lua.SyntaxKind.SubtractionOperator,
                 node
             );
+        }
+
         case "substr":
-            if (node.arguments.length === 1) {
-                const argument = context.transformExpression(node.arguments[0]);
-                const arg1 = createExpressionPlusOne(argument);
-                return createStringCall("sub", node, caller, arg1);
-            } else {
-                const sumArg = lua.createBinaryExpression(params[0], params[1], lua.SyntaxKind.AdditionOperator);
-                return createStringCall("sub", node, caller, createExpressionPlusOne(params[0]), sumArg);
-            }
+            return transformLuaLibFunction(context, LuaLibFeature.StringSubstr, node, caller, ...params);
         case "substring":
-            if (node.arguments.length === 1) {
-                const arg1 = createExpressionPlusOne(params[0]);
-                return createStringCall("sub", node, caller, arg1);
-            } else {
-                const arg1 = createExpressionPlusOne(params[0]);
-                const arg2 = params[1];
-                return createStringCall("sub", node, caller, arg1, arg2);
-            }
-        case "slice":
-            if (node.arguments.length === 0) {
-                return caller;
-            } else if (node.arguments.length === 1) {
-                const arg1 = createExpressionPlusOne(params[0]);
-                return createStringCall("sub", node, caller, arg1);
-            } else {
-                const arg1 = createExpressionPlusOne(params[0]);
-                const arg2 = params[1];
-                return createStringCall("sub", node, caller, arg1, arg2);
+            return transformLuaLibFunction(context, LuaLibFeature.StringSubstring, node, caller, ...params);
+
+        case "slice": {
+            const literalArg1 = getNumberLiteralValue(params[0]);
+            if (params[0] && literalArg1 !== undefined) {
+                let stringSubArgs: lua.Expression[] | undefined = [
+                    addToNumericExpression(params[0], literalArg1 < 0 ? 0 : 1),
+                ];
+
+                if (params[1]) {
+                    const literalArg2 = getNumberLiteralValue(params[1]);
+                    if (literalArg2 !== undefined) {
+                        stringSubArgs.push(addToNumericExpression(params[1], literalArg2 < 0 ? -1 : 0));
+                    } else {
+                        stringSubArgs = undefined;
+                    }
+                }
+
+                // Inline string.sub call if we know that both parameters are pure and aren't negative
+                if (stringSubArgs) {
+                    return createStringCall("sub", node, caller, ...stringSubArgs);
+                }
             }
+
+            return transformLuaLibFunction(context, LuaLibFeature.StringSlice, node, caller, ...params);
+        }
+
         case "toLowerCase":
             return createStringCall("lower", node, caller);
         case "toUpperCase":
@@ -89,13 +99,32 @@ export function transformStringPrototypeCall(
             return transformLuaLibFunction(context, LuaLibFeature.StringTrimStart, node, caller);
         case "split":
             return transformLuaLibFunction(context, LuaLibFeature.StringSplit, node, caller, ...params);
-        case "charAt":
-            const firstParamPlusOne = createExpressionPlusOne(params[0]);
-            return createStringCall("sub", node, caller, firstParamPlusOne, firstParamPlusOne);
+
+        case "charAt": {
+            const literalValue = getNumberLiteralValue(params[0]);
+            // Inline string.sub call if we know that parameter is pure and isn't negative
+            if (literalValue !== undefined && literalValue >= 0) {
+                const firstParamPlusOne = addToNumericExpression(params[0], 1);
+                return createStringCall("sub", node, caller, firstParamPlusOne, firstParamPlusOne);
+            }
+
+            return transformLuaLibFunction(context, LuaLibFeature.StringCharAt, node, caller, ...params);
+        }
+
         case "charCodeAt": {
-            const firstParamPlusOne = createExpressionPlusOne(params[0]);
-            return createStringCall("byte", node, caller, firstParamPlusOne);
+            const literalValue = getNumberLiteralValue(params[0]);
+            // Inline string.sub call if we know that parameter is pure and isn't negative
+            if (literalValue !== undefined && literalValue >= 0) {
+                return lua.createBinaryExpression(
+                    createStringCall("byte", node, caller, addToNumericExpression(params[0], 1)),
+                    createNaN(),
+                    lua.SyntaxKind.OrOperator
+                );
+            }
+
+            return transformLuaLibFunction(context, LuaLibFeature.StringCharCodeAt, node, caller, ...params);
         }
+
         case "startsWith":
             return transformLuaLibFunction(context, LuaLibFeature.StringStartsWith, node, caller, ...params);
         case "endsWith":
diff --git a/src/transformation/utils/lua-ast.ts b/src/transformation/utils/lua-ast.ts
@@ -22,24 +22,45 @@ export function createExportsIdentifier(): lua.Identifier {
     return lua.createIdentifier("____exports");
 }
 
-export function createExpressionPlusOne(expression: lua.Expression): lua.Expression {
-    if (lua.isNumericLiteral(expression)) {
-        const newNode = lua.cloneNode(expression);
-        newNode.value += 1;
+export function addToNumericExpression(expression: lua.Expression, change: number): lua.Expression {
+    if (change === 0) return expression;
+
+    const literalValue = getNumberLiteralValue(expression);
+    if (literalValue !== undefined) {
+        const newNode = lua.createNumericLiteral(literalValue + change);
+        lua.setNodePosition(newNode, expression);
         return newNode;
     }
 
     if (lua.isBinaryExpression(expression)) {
         if (
-            expression.operator === lua.SyntaxKind.SubtractionOperator &&
             lua.isNumericLiteral(expression.right) &&
-            expression.right.value === 1
+            ((expression.operator === lua.SyntaxKind.SubtractionOperator && expression.right.value === change) ||
+                (expression.operator === lua.SyntaxKind.AdditionOperator && expression.right.value === -change))
         ) {
             return expression.left;
         }
     }
 
-    return lua.createBinaryExpression(expression, lua.createNumericLiteral(1), lua.SyntaxKind.AdditionOperator);
+    return change > 0
+        ? lua.createBinaryExpression(expression, lua.createNumericLiteral(change), lua.SyntaxKind.AdditionOperator)
+        : lua.createBinaryExpression(expression, lua.createNumericLiteral(-change), lua.SyntaxKind.SubtractionOperator);
+}
+
+export function getNumberLiteralValue(expression?: lua.Expression) {
+    if (!expression) return undefined;
+
+    if (lua.isNumericLiteral(expression)) return expression.value;
+
+    if (
+        lua.isUnaryExpression(expression) &&
+        expression.operator === lua.SyntaxKind.NegationOperator &&
+        lua.isNumericLiteral(expression.operand)
+    ) {
+        return -expression.operand.value;
+    }
+
+    return undefined;
 }
 
 export function createImmediatelyInvokedFunctionExpression(
@@ -197,3 +218,11 @@ export function createLocalOrExportedOrGlobalDeclaration(
         return [];
     }
 }
+
+export const createNaN = (tsOriginal?: ts.Node) =>
+    lua.createBinaryExpression(
+        lua.createNumericLiteral(0),
+        lua.createNumericLiteral(0),
+        lua.SyntaxKind.DivisionOperator,
+        tsOriginal
+    );
diff --git a/src/transformation/visitors/access.ts b/src/transformation/visitors/access.ts
diff --git a/test/unit/builtins/array.spec.ts b/test/unit/builtins/array.spec.ts
diff --git a/test/unit/builtins/string.spec.ts b/test/unit/builtins/string.spec.ts

Original file line number	Diff line number	Diff line change
`@@ -13,5 +13,5 @@ function __TS__StringPadEnd(this: string, maxLength: number, fillString = " "):`
`13`	`13`	`fillString += fillString.repeat(maxLength / fillString.length);`
`14`	`14`	`}`
`15`	`15`
`16`		`- return this + fillString.slice(0, Math.floor(maxLength));`
	`16`	`+ return this + string.sub(fillString, 1, Math.floor(maxLength));`
`17`	`17`	`}`