diff --git a/src/Unluau/Lifter/Expressions/BinaryExpression.cs b/src/Unluau/Lifter/Expressions/BinaryExpression.cs index 0456919..a810b9f 100644 --- a/src/Unluau/Lifter/Expressions/BinaryExpression.cs +++ b/src/Unluau/Lifter/Expressions/BinaryExpression.cs @@ -39,6 +39,13 @@ public BinaryExpression(Expression left, BinaryOperation operation, Expression r public override void Write(Output output) { + // Add null checks + if (Left == null || Right == null) + { + output.Write("nil"); + return; + } + Left.Write(output); output.Write($" {BinaryOperationChar(Operation)} "); Right.Write(output); diff --git a/src/Unluau/Lifter/Expressions/Closure.cs b/src/Unluau/Lifter/Expressions/Closure.cs index f55899f..7e1c3e2 100644 --- a/src/Unluau/Lifter/Expressions/Closure.cs +++ b/src/Unluau/Lifter/Expressions/Closure.cs @@ -44,7 +44,11 @@ public override void Write(Output output) output.WriteLine(")"); - Block.Write(output); + // Add null check for Block + if (Block != null) + Block.Write(output); + else + output.WriteLine("-- [Unluau: Missing function body]"); output.Write("end"); } diff --git a/src/Unluau/Lifter/Expressions/NameIndex.cs b/src/Unluau/Lifter/Expressions/NameIndex.cs index 55f359e..3982fde 100644 --- a/src/Unluau/Lifter/Expressions/NameIndex.cs +++ b/src/Unluau/Lifter/Expressions/NameIndex.cs @@ -24,6 +24,13 @@ public NameIndex(Expression expression, string name, bool isSelf = false) public override void Write(Output output) { + // Add null check + if (Expression == null) + { + output.Write("nil"); + return; + } + Expression.Write(output); output.Write((IsSelf ? ":" : ".") + Name); } diff --git a/src/Unluau/Lifter/Lifter.cs b/src/Unluau/Lifter/Lifter.cs index b0ef726..42a85d5 100644 --- a/src/Unluau/Lifter/Lifter.cs +++ b/src/Unluau/Lifter/Lifter.cs @@ -125,15 +125,23 @@ private Block LiftBlock(Function function, Registers registers, int pcStart = 0, if (options.PerferStringInterpolation) { - // If we perfer string interpolation set the expression as an interpolated string - string format = (((LocalExpression)nameIndex!.Expression).Expression as StringLiteral)!.Value; - expression = new InterpolatedStringLiteral(format, new List()); + // Safely extract the format string + var localExpr = nameIndex?.Expression as LocalExpression; + var stringLiteral = localExpr?.Expression as StringLiteral; + if (stringLiteral != null) + { + string format = stringLiteral.Value; + expression = new InterpolatedStringLiteral(format, new List()); + } } else { // Otherwise add an expression group around the string literal - nameIndex!.Expression = new ExpressionGroup(nameIndex!.Expression); - expression = nameIndex; + if (nameIndex != null) + { + nameIndex.Expression = new ExpressionGroup(nameIndex.Expression); + expression = nameIndex; + } } } @@ -293,7 +301,13 @@ private Block LiftBlock(Function function, Registers registers, int pcStart = 0, case OpCode.SETTABLEKS: { StringConstant target = (StringConstant)function.GetConstant(++pc); - Expression table = registers.GetExpression(instruction.B), tableValue = ((LocalExpression)table).Expression; + Expression table = registers.GetExpression(instruction.B); + if (table == null || !(table is LocalExpression)) + { + // Skip this instruction if table is null or not a LocalExpression + break; + } + Expression tableValue = ((LocalExpression)table).Expression; if (options.InlineTableDefintions && tableValue is TableLiteral) { @@ -317,7 +331,12 @@ private Block LiftBlock(Function function, Registers registers, int pcStart = 0, case OpCode.SETTABLE: { Expression expression = registers.GetRefExpressionValue(instruction.C), value = registers.GetExpression(instruction.A); - Expression table = registers.GetExpression(instruction.B, false), tableValue = ((LocalExpression)table).Expression; + Expression table = registers.GetExpression(instruction.B, false); + if (table == null || !(table is LocalExpression)) + { + break; + } + Expression tableValue = ((LocalExpression)table).Expression; if (options.InlineTableDefintions && tableValue is TableLiteral) { @@ -682,7 +701,11 @@ private Block LiftBlock(Function function, Registers registers, int pcStart = 0, private int IsSelf(Expression expression) { - Expression? value = expression.GetValue(); + Expression? value = expression?.GetValue(); + + // Add null check here + if (value == null) + return 0; if (value is NameIndex nameIndex) return nameIndex.IsSelf ? 1 : 0; diff --git a/src/Unluau/Lifter/Statements/Assignment.cs b/src/Unluau/Lifter/Statements/Assignment.cs index 9932724..24864ee 100644 --- a/src/Unluau/Lifter/Statements/Assignment.cs +++ b/src/Unluau/Lifter/Statements/Assignment.cs @@ -22,6 +22,13 @@ public Assignment(Expression variable, Expression value) public override void Write(Output output) { + // Add null checks + if (Variable == null || Value == null) + { + output.Write("-- [Unluau: Failed to decompile assignment]"); + return; + } + if (Value is LocalExpression && ((LocalExpression)Value).Expression is Closure) { output.Write("function "); diff --git a/src/Unluau/Lifter/Statements/IfElse.cs b/src/Unluau/Lifter/Statements/IfElse.cs index cf037bb..b3d2fff 100644 --- a/src/Unluau/Lifter/Statements/IfElse.cs +++ b/src/Unluau/Lifter/Statements/IfElse.cs @@ -24,6 +24,13 @@ public IfElse(Expression condition, Block ifBody) public override void Write(Output output) { + // Add null checks for Condition and IfBody + if (Condition == null || IfBody == null) + { + output.WriteLine("-- [Unluau: Failed to decompile if statement - missing condition or body]"); + return; + } + output.Write("if "); Condition.Write(output); output.WriteLine(" then"); diff --git a/src/Unluau/Lifter/Statements/LocalAssignment.cs b/src/Unluau/Lifter/Statements/LocalAssignment.cs index bfa05eb..5bed08e 100644 --- a/src/Unluau/Lifter/Statements/LocalAssignment.cs +++ b/src/Unluau/Lifter/Statements/LocalAssignment.cs @@ -52,6 +52,13 @@ public bool TryGetVariables(out ExpressionList variable) public override void Write(Output output) { + // Add null checks + if (Expression == null || Value == null) + { + output.Write("-- [Unluau: Failed to decompile local assignment]"); + return; + } + output.Write("local "); if (Value is Closure) diff --git a/src/Unluau/Lifter/Statements/Return.cs b/src/Unluau/Lifter/Statements/Return.cs index b494eda..b4a7dd3 100644 --- a/src/Unluau/Lifter/Statements/Return.cs +++ b/src/Unluau/Lifter/Statements/Return.cs @@ -22,21 +22,25 @@ public Return() public override void Write(Output output) { - output.Write("return "); - - bool first = true; - - foreach (Expression expression in Expressions) + output.Write("return"); + + if (Expressions != null && Expressions.Count > 0) { - if (expression is null) - continue; - - expression.Write(output); - - if (first) + output.Write(" "); + bool first = true; + + foreach (Expression expression in Expressions) + { + if (!first) + output.Write(", "); + + if (expression == null) + output.Write("nil"); + else + expression.Write(output); + first = false; - else - output.Write(","); + } } } } diff --git a/src/Unluau/Lifter/Statements/Statement.cs b/src/Unluau/Lifter/Statements/Statement.cs index 20caa83..4f16de3 100644 --- a/src/Unluau/Lifter/Statements/Statement.cs +++ b/src/Unluau/Lifter/Statements/Statement.cs @@ -14,6 +14,14 @@ public static void WriteSequence(Output output, IList statements) for (int i = 0; i < statements.Count; i++) { var statement = statements[i]; + + // Add null check for statement + if (statement == null) + { + output.WriteLine("-- [Unluau: null statement]"); + continue; + } + if (statement.Comment != null) output.WriteLine("-- " + statement.Comment);