Proper implementation of local var declarations

Underanalyzer/Decompiler/AST/ASTFragmentContext.cs

@@ -65,6 +65,16 @@ public sealed class ASTFragmentContext
     /// </summary>
     public List<string> LocalVariableNamesList { get; } = [];
 
+    /// <summary>
+    /// The current local variable scope.
+    /// </summary>
+    internal LocalScope? CurrentLocalScope { get; private set; }
+
+    /// <summary>
+    /// The current block being cleaned up during post-cleanup.
+    /// </summary>
+    internal BlockNode? CurrentPostCleanupBlock { get; set; }
+
     /// <summary>
     /// Map of code entry names to function names, for all children fragments/sub-functions of this context.
     /// </summary>
@@ -117,9 +127,8 @@ internal ASTFragmentContext(Fragment fragment)
     /// </summary>
     internal void RemoveLocal(string name)
     {
-        if (LocalVariableNames.Contains(name))
+        if (LocalVariableNames.Remove(name))
         {
-            LocalVariableNames.Remove(name);
             LocalVariableNamesList.Remove(name);
         }
     }
@@ -166,4 +175,57 @@ internal void RemoveLocal(string name)
         NamedArgumentByIndex[index] = name;
         return name;
     }
+
+    /// <summary>
+    /// Pushes a new local scope for this fragment.
+    /// </summary>
+    internal void PushLocalScope(DecompileContext context, BlockNode containingBlock, IStatementNode startStatement)
+    {
+        if (!context.Settings.CleanupLocalVarDeclarations)
+        {
+            return;
+        }
+
+        LocalScope? oldScope = CurrentLocalScope;
+        CurrentLocalScope = new LocalScope(oldScope, containingBlock, startStatement);
+        oldScope?.Children?.Add(CurrentLocalScope);
+    }
+
+    /// <summary>
+    /// Pops a local scope for this fragment.
+    /// </summary>
+    internal void PopLocalScope(DecompileContext context)
+    {
+        if (!context.Settings.CleanupLocalVarDeclarations)
+        {
+            return;
+        }
+
+        LocalScope oldScope = CurrentLocalScope!;
+        CurrentLocalScope = oldScope.Parent;
+
+        if (CurrentLocalScope is null)
+        {
+            // Perform a pass on local scopes, generating local variable declarations.
+            HashSet<string> alreadyDeclared = new(LocalVariableNames.Count);
+            HashSet<string> declaredAnywhere = new(LocalVariableNames.Count);
+            oldScope.GenerateDeclarations(alreadyDeclared, declaredAnywhere);
+
+            // Generate local variable declaration for any remaining undeclared locals
+            List<string> toDeclareAtTop = new(LocalVariableNames.Count);
+            foreach (string local in LocalVariableNamesList)
+            {
+                if (!declaredAnywhere.Contains(local))
+                {
+                    toDeclareAtTop.Add(local);
+                }
+            }
+            if (toDeclareAtTop.Count > 0)
+            {
+                LocalVarDeclNode decl = new();
+                decl.Locals.AddRange(toDeclareAtTop);
+                oldScope.ContainingBlock.Children.Insert(0, decl);
+            }
+        }
+    }
 }

Underanalyzer/Decompiler/AST/IASTNode.cs

@@ -17,6 +17,12 @@ public interface IASTNode<T>
     /// </summary>
     public T Clean(ASTCleaner cleaner);
 
+    /// <summary>
+    /// Second cleanup pass on this node and all of its sub-nodes.
+    /// Returns the second-pass cleaned version of the node (which is very often the same reference).
+    /// </summary>
+    public T PostClean(ASTCleaner cleaner);
+
     /// <summary>
     /// Prints this node using the provided printer.
     /// </summary>

Underanalyzer/Decompiler/AST/IStatementNode.cs

@@ -19,10 +19,10 @@ public interface IStatementNode : IASTNode<IStatementNode>
     /// <summary>
     /// If true, an empty line should be printed before this node, unless at the start of a block.
     /// </summary>
-    public bool EmptyLineBefore { get; }
+    public bool EmptyLineBefore { get; set; }
 
     /// <summary>
     /// If true, an empty line should be printed after this node, unless at the end of a block.
     /// </summary>
-    public bool EmptyLineAfter { get; }
+    public bool EmptyLineAfter { get; set; }
 }

Underanalyzer/Decompiler/AST/LocalScope.cs

@@ -0,0 +1,224 @@
+/*
+  This Source Code Form is subject to the terms of the Mozilla Public
+  License, v. 2.0. If a copy of the MPL was not distributed with this
+  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+*/
+
+using System.Collections.Generic;
+
+namespace Underanalyzer.Decompiler.AST;
+
+/// <summary>
+/// Represents a scope within the AST for which local variables can be declared.
+/// </summary>
+internal sealed class LocalScope(LocalScope? parent, BlockNode containingBlock, IStatementNode startStatement)
+{
+    /// <summary>
+    /// The block that contains this local scope, directly.
+    /// This block must directly contain <see cref="StartStatement"/>.
+    /// </summary>
+    public BlockNode ContainingBlock { get; } = containingBlock;
+
+    /// <summary>
+    /// The first statement that is within this local scope.
+    /// For example, with an if statement's true/false (else) blocks, this 
+    /// would be the enveloping overall <see cref="IfNode"/>.
+    /// </summary>
+    /// <remarks>
+    /// Specifically for blocks at the root of their fragment, 
+    /// this is the same as <see cref="ContainingBlock"/>.
+    /// </remarks>
+    public IStatementNode StartStatement { get; } = startStatement;
+
+    /// <summary>
+    /// Scopes contained within this local scope.
+    /// </summary>
+    public List<LocalScope> Children { get; } = new(4);
+
+    /// <summary>
+    /// Scope containing this local scope, if one exists; null otherwise.
+    /// </summary>
+    public LocalScope? Parent { get; private set; } = parent;
+
+    /// <summary>
+    /// Local variable names that have been denoted as "declared" for this local scope.
+    /// This means that, each local variable, if not declared already, should be declared
+    /// before its first usage inside of the scope.
+    /// </summary>
+    public HashSet<string> DeclaredLocals { get; } = new(4);
+
+    /// <summary>
+    /// For local variables that are denoted as "declared" for this scope,
+    /// this is a map of their names to their first <see cref="AssignNode"/>,
+    /// for use in generating declarations.
+    /// </summary>
+    public Dictionary<string, AssignNode> FirstLocalAssignments { get; } = new(4);
+
+    /// <summary>
+    /// Local variable names that were specifically hoisted to be declared just 
+    /// *before* this local scope's <see cref="StartStatement"/>.
+    /// </summary>
+    public List<string> HoistedLocals { get; } = new(4);
+
+    // Prevents scope from being considered too many times during certain search operations
+    private bool _ignoreInSearch = false;
+
+    /// <summary>
+    /// Inserts this local scope as a child of another existing local scope.
+    /// </summary>
+    public void InsertAsChildOf(LocalScope parent)
+    {
+        Parent = parent;
+        parent.Children.Add(this);
+    }
+
+    /// <summary>
+    /// Returns whether a local variable has already been denoted as "declared" in this 
+    /// local scope, or any of its parent local scopes.
+    /// </summary>
+    public bool LocalDeclaredInAnyParentOrSelf(string localName)
+    {
+        if (DeclaredLocals.Contains(localName))
+        {
+            return true;
+        }
+        if (Parent is not null)
+        {
+            return Parent.LocalDeclaredInAnyParentOrSelf(localName);
+        }
+        return false;
+    }
+
+    /// <summary>
+    /// Enumerates all children scopes in order, finding the first one that either itself
+    /// declares the given local variable name, or one of its own children declares the given
+    /// local variable name. If no such child exists, returns null.
+    /// 
+    /// Also returns the precise local scope where the local is currently declared, or null if none.
+    /// </summary>
+    public (LocalScope?, LocalScope?) FindLocalDeclaredInAnyChild(string localName)
+    {
+        foreach (LocalScope child in Children)
+        {
+            if (child._ignoreInSearch)
+            {
+                // Already considered and ruled out this child, no need to do it again
+                continue;
+            }
+            if (child.DeclaredLocals.Contains(localName))
+            {
+                return (child, child);
+            }
+            if (child.FindLocalDeclaredInAnyChild(localName) is (LocalScope, LocalScope declaration))
+            {
+                return (child, declaration);
+            }
+        }
+        return (null, null);
+    }
+
+    /// <summary>
+    /// In the situation that a local variable is not marked as declared in a scope, 
+    /// nor its parents (see <see cref="LocalDeclaredInAnyParentOrSelf(string)"/>),
+    /// this can be used to return the best local scope to hoist an earlier local variable
+    /// declaration to. Specifically, it should be hoisted to just before the scope.
+    /// 
+    /// Returns null if no suitable scope is found (i.e., the local is never declared anywhere).
+    /// Also returns the precise local scope where the local is currently declared, or null if none.
+    /// </summary>
+    public (LocalScope?, LocalScope?) FindBestHoistLocation(string localName)
+    {
+        // Look for locals in any immediate child scopes first
+        if (FindLocalDeclaredInAnyChild(localName) is (LocalScope result, LocalScope declaration))
+        {
+            return (result, declaration);
+        }
+
+        // Since none were found here, search the parent, if one exists
+        if (Parent is not null)
+        {
+            // Prevent this scope from being searched for locals (we already considered it)
+            _ignoreInSearch = true;
+
+            (LocalScope?, LocalScope?) parentResult = Parent.FindBestHoistLocation(localName);
+
+            _ignoreInSearch = false;
+            return parentResult;
+        }
+
+        // No parent, so no result
+        return (null, null);
+    }
+
+    /// <summary>
+    /// Generates local variable declarations where necessary.
+    /// Recursively performs this operation across all local scopes in the fragment.
+    /// </summary>
+    public void GenerateDeclarations(HashSet<string> alreadyDeclared, HashSet<string> declaredAnywhere)
+    {
+        // First, hoist local variables from child scopes.
+        foreach (LocalScope child in Children)
+        {
+            if (child.HoistedLocals.Count > 0)
+            {
+                // This child scope has hoisted locals declared before its scope.
+
+                // Generate a local var declaration, and add those locals to it.
+                LocalVarDeclNode? localDecl = null;
+                bool anyNew = false;
+                foreach (string hoistedLocal in child.HoistedLocals)
+                {
+                    if (alreadyDeclared.Add(hoistedLocal))
+                    {
+                        anyNew = true;
+                        localDecl ??= new();
+                        localDecl.Locals.Add(hoistedLocal);
+                        declaredAnywhere.Add(hoistedLocal);
+                    }
+                }
+
+                // Ensure there's at least one local being hoisted
+                if (!anyNew)
+                {
+                    continue;
+                }
+
+                // Find start statement and insert just before it
+                List<IStatementNode> blockChildren = child.ContainingBlock.Children;
+                int index = blockChildren.IndexOf(child.StartStatement);
+                if (index == -1)
+                {
+                    // Failsafe; start at 0. This usually shouldn't happen, though.
+                    index = 0;
+                }
+                blockChildren.Insert(index, localDecl!);
+
+                // If immediately before a node that has an empty line before it, take it over.
+                if (index < blockChildren.Count - 1 && blockChildren[index + 1].EmptyLineBefore)
+                {
+                    blockChildren[index + 1].EmptyLineBefore = false;
+                    localDecl!.EmptyLineBefore = true;
+                }
+            }
+        }
+
+        // For remaining locals declared in this scope, declare them on their first assignments
+        foreach (string local in DeclaredLocals)
+        {
+            if (alreadyDeclared.Add(local))
+            {
+                FirstLocalAssignments[local].DeclareLocalVar = true;
+                declaredAnywhere.Add(local);
+            }
+        }
+
+        // Generate declarations for child scopes
+        foreach (LocalScope child in Children)
+        {
+            child.GenerateDeclarations(alreadyDeclared, declaredAnywhere);
+        }
+
+        // Exit this scope (un-declare its declared locals)
+        alreadyDeclared.ExceptWith(DeclaredLocals);
+    }
+}

Underanalyzer/Decompiler/AST/Nodes/ArrayInitNode.cs

@@ -35,6 +35,15 @@ public IExpressionNode Clean(ASTCleaner cleaner)
         return this;
     }
 
+    public IExpressionNode PostClean(ASTCleaner cleaner)
+    {
+        for (int i = 0; i < Elements.Count; i++)
+        {
+            Elements[i] = Elements[i].PostClean(cleaner);
+        }
+        return this;
+    }
+
     public void Print(ASTPrinter printer)
     {
         printer.Write('[');

Underanalyzer/Decompiler/AST/Nodes/AssetReferenceNode.cs

@@ -35,6 +35,11 @@ public IExpressionNode Clean(ASTCleaner cleaner)
         return this;
     }
 
+    public IExpressionNode PostClean(ASTCleaner cleaner)
+    {
+        return this;
+    }
+
     public bool RequiresMultipleLines(ASTPrinter printer)
     {
         return false;