CodeGeneratorVisitor.cs

using Antlr4.Runtime.Tree;
using System;
using System.Collections.Generic;

namespace TinyLanguage
{
    // OpCode 枚举定义在 CPU.cs 中，并在此处通过命名空间直接使用。
    // 如果 OpCode 定义在其他地方，需要确保它在此处是可访问的。

    /// <summary>
    /// 符号表条目，用于存储变量信息。
    /// </summary>
    public class SymbolTableEntry
    {
        /// <summary>
        /// 变量在VM内存中的地址。
        /// </summary>
        public ushort Address { get; set; }
        // public string Type { get; set; } // 当前 TL 语言只有 INT 类型，所以类型字段暂时省略
    }

    /// <summary>
    /// ANTLR Visitor 类，用于遍历解析树并生成目标VM的字节码。
    /// 继承自 ANTLR 生成的 TLBaseVisitor<object>。
    /// </summary>
    public class CodeGeneratorVisitor : TinyLanguageBaseVisitor<object> // object 是默认的 Visit 方法返回值类型
    {
        // 存储生成的字节码的列表
        private readonly List<byte> bytecode = new List<byte>();
        // 符号表，用于存储变量名及其内存地址
        private readonly Dictionary<string, SymbolTableEntry> symbolTable = new Dictionary<string, SymbolTableEntry>();
        //下一个可用的全局变量内存地址，从0x0200开始（可根据需要调整）
        private ushort nextMemoryAddress = 0x0200;

        // --- 寄存器约定 (与VM设计一致) ---
        private const byte REG_ACCUM = 5;         // R5: 表达式计算的累加器/主要结果寄存器
        private const byte REG_TEMP = 6;          // R6: 用于二元运算的临时操作数寄存器
        private const byte REG_COND_RESULT = 4;   // R4: 存储条件判断的布尔结果 (0 表示 false, 1 表示 true)
        private const byte REG_CMP_TEMP = 3;      // R3: 用于比较操作时存储中间差值 (例如 a-b)

        /// <summary>
        /// 获取已生成的字节码列表的副本。
        /// </summary>
        /// <returns>字节码列表</returns>
        public List<byte> GetBytecode()
        {
            return new List<byte>(bytecode); // 返回副本，防止外部直接修改内部列表
        }

        // --- 辅助方法：发射字节码 ---

        /// <summary>
        /// 向字节码列表中添加一个字节。
        /// </summary>
        private void EmitByte(byte val)
        {
            bytecode.Add(val);
        }

        /// <summary>
        /// 将 OpCode 枚举成员的值（即其字节表示）添加到字节码列表。
        /// </summary>
        private void EmitOpCode(OpCode op)
        {
            EmitByte((byte)op); // 将 OpCode 强制转换为 byte
        }

        /// <summary>
        /// 将一个16位有符号短整型 (short) 以小端字节序发射到字节码列表 (2字节)。
        /// 通常用于立即数。
        /// </summary>
        private void EmitShort(short value)
        {
            EmitByte((byte)(value & 0xFF));          // 低字节
            EmitByte((byte)((value >> 8) & 0xFF));   // 高字节
        }

        /// <summary>
        /// 将一个16位无符号短整型 (ushort) 以小端字节序发射到字节码列表 (2字节)。
        /// 通常用于内存地址。
        /// </summary>
        private void EmitUShort(ushort value)
        {
            EmitByte((byte)(value & 0xFF));          // 低字节
            EmitByte((byte)((value >> 8) & 0xFF));   // 高字节
        }

        // --- 辅助方法：发射特定指令 ---

        private void EmitConst(byte reg, short value)
        {
            EmitOpCode(OpCode.CONST);
            EmitByte(reg);
            EmitShort(value);
        }

        private void EmitLoad(byte reg, ushort memAddr)
        {
            EmitOpCode(OpCode.LOAD);
            EmitByte(reg);
            EmitUShort(memAddr);
        }

        private void EmitStore(byte reg, ushort memAddr)
        {
            EmitOpCode(OpCode.STORE);
            EmitByte(reg);
            EmitUShort(memAddr);
        }

        private void EmitMov(byte rd, byte rs)
        {
            EmitOpCode(OpCode.MOV);
            EmitByte(rd);
            EmitByte(rs);
        }

        /// <summary>
        /// 发射需要三个寄存器操作数的算术或逻辑指令 (如 ADD, SUB, MUL, DIV)。
        /// </summary>
        private void EmitOpRegs(OpCode op, byte rd, byte r1, byte r2)
        {
            EmitOpCode(op);
            EmitByte(rd);
            EmitByte(r1);
            EmitByte(r2);
        }

        private void EmitPush(byte reg)
        {
            EmitOpCode(OpCode.PUSH);
            EmitByte(reg);
        }

        private void EmitPop(byte reg)
        {
            EmitOpCode(OpCode.POP);
            EmitByte(reg);
        }

        // --- 标签和跳转管理 ---

        /// <summary>
        /// 创建一个新的列表，用于存储需要回填地址的跳转指令位置。
        /// </summary>
        private List<int> CreateLabelPlaceholderList()
        {
            return new List<int>();
        }

        /// <summary>
        /// 发射一个无条件跳转指令 (JMP)，其目标地址暂时为占位符。
        /// 将占位符地址的字节码索引添加到 patchListForThisJump 列表中，以便后续回填。
        /// </summary>
        private void EmitJmp(List<int> patchListForThisJump)
        {
            EmitOpCode(OpCode.JMP);
            patchListForThisJump.Add(bytecode.Count); // 记录地址的低字节在字节码中的索引
            EmitUShort(0xFFFF); // 发射一个16位占位符地址 (例如 0xFFFF)
        }

        /// <summary>
        /// 发射一个条件跳转指令 (如 JZ, JL)，其目标地址暂时为占位符。
        /// 将占位符地址的字节码索引添加到 patchListForThisJump 列表中。
        /// </summary>
        private void EmitCondJmp(OpCode op, byte regTest, List<int> patchListForThisJump)
        {
            EmitOpCode(op);
            EmitByte(regTest);
            patchListForThisJump.Add(bytecode.Count); // 记录地址的低字节在字节码中的索引
            EmitUShort(0xFFFF); // 发射一个16位占位符地址
        }

        /// <summary>
        /// 将当前字节码位置作为目标地址，回填到 patchList 中记录的所有跳转指令的占位符地址处。
        /// </summary>
        private void PlaceLabel(List<int> patchList)
        {
            ushort targetAddress = (ushort)bytecode.Count; // 当前字节码长度即为标签的目标地址
            foreach (int patchLocation_LowByteIndex in patchList)
            {
                bytecode[patchLocation_LowByteIndex] = (byte)(targetAddress & 0xFF);        // 回填低字节
                bytecode[patchLocation_LowByteIndex + 1] = (byte)((targetAddress >> 8) & 0xFF); // 回填高字节
            }
            patchList.Clear(); // 清空已处理的占位符列表 (可选，但有助于避免重复使用)
        }

        // --- Visitor 方法实现 (重写 TLBaseVisitor 中的方法) ---

        /// <summary>
        /// 访问程序根节点 (program)。
        /// </summary>
        public override object VisitProgram(TinyLanguageParser.ProgramContext context)
        {
            Visit(context.statement_list()); // 访问程序主体语句列表
            // HLT 指令应在调用此 Visitor 完成后，由主编译器逻辑添加到字节码的末尾。
            return null; // Visitor 方法通常返回 null 或一个通用对象
        }

        /// <summary>
        /// 访问语句列表节点 (statement_list)。
        /// </summary>
        public override object VisitStatement_list(TinyLanguageParser.Statement_listContext context)
        {
            foreach (var stmtCtx in context.statement()) // 遍历列表中的每一个语句
            {
                Visit(stmtCtx); // 访问该语句
            }
            return null;
        }

        // VisitStatement 方法通常不需要重写，ANTLR的机制会自动将其分发到
        // 更具体的语句类型访问方法 (如 VisitAssignment, VisitIf_statement 等)。
        // public override object VisitStatement(TinyLanguageParser.StatementContext context)
        // {
        //     return base.VisitStatement(context); // 或者 VisitChildren(context)
        // }

        /// <summary>
        /// 访问变量声明节点 (declaration)。
        /// </summary>
        public override object VisitDeclaration(TinyLanguageParser.DeclarationContext context)
        {
            string varName = context.IDENTIFIER().GetText(); // 获取变量名
            if (symbolTable.ContainsKey(varName))
            {
                // 可以通过抛出异常或记录错误信息来处理重复声明
                throw new Exception($"语义错误: 变量 '{varName}' 在符号表中已存在，发生重复声明。");
            }

            // 将变量添加到符号表，并分配内存地址
            symbolTable[varName] = new SymbolTableEntry { Address = nextMemoryAddress };
            nextMemoryAddress += 2; // INT 类型在VM中占用2个字节
            return null; // 变量声明本身不生成执行时字节码
        }

        /// <summary>
        /// 访问赋值语句节点 (assignment)。
        /// </summary>
        public override object VisitAssignment(TinyLanguageParser.AssignmentContext context)
        {
            string varName = context.IDENTIFIER().GetText(); // 获取被赋值的变量名
            if (!symbolTable.TryGetValue(varName, out SymbolTableEntry entry))
            {
                throw new Exception($"语义错误: 变量 '{varName}' 未声明，不能进行赋值。");
            }

            Visit(context.expression()); // 访问并计算右侧表达式的值，结果存入 REG_ACCUM (R5)
            EmitStore(REG_ACCUM, entry.Address); // 将 R5 中的结果存储到变量的内存地址
            return null;
        }

        /// <summary>
        /// 访问输入语句节点 (input_statement)。
        /// </summary>
        public override object VisitInput_statement(TinyLanguageParser.Input_statementContext context)
        {
            string varName = context.IDENTIFIER().GetText();
            if (!symbolTable.TryGetValue(varName, out SymbolTableEntry entry))
            {
                throw new Exception($"语义错误: 变量 '{varName}' 未声明，不能用于接收输入。");
            }

            EmitOpCode(OpCode.IN);        // 发射 IN 指令
            EmitByte(REG_ACCUM);          // 指定输入结果存入 REG_ACCUM (R5)
            EmitStore(REG_ACCUM, entry.Address); // 将 R5 中的输入结果存储到变量的内存地址
            return null;
        }

        /// <summary>
        /// 访问输出语句节点 (output_statement)。
        /// </summary>
        public override object VisitOutput_statement(TinyLanguageParser.Output_statementContext context)
        {
            Visit(context.expression()); // 访问并计算表达式的值，结果存入 REG_ACCUM (R5)
            EmitOpCode(OpCode.OUT);       // 发射 OUT 指令
            EmitByte(REG_ACCUM);          // 指定输出 REG_ACCUM (R5) 中的值
            return null;
        }

        // --- 表达式求值 ---
        // 约定：所有表达式求值后，其结果都应存放在 REG_ACCUM (R5) 寄存器中。

        /// <summary>
        /// 访问乘除法表达式节点 (MulDivExpr)。
        /// </summary>
        public override object VisitMulDivExpr(TinyLanguageParser.MulDivExprContext context)
        {
            Visit(context.expression(0));   // 计算左操作数，结果在 R5
            EmitPush(REG_ACCUM);            // 将左操作数的值压入VM栈中保存
            Visit(context.expression(1));   // 计算右操作数，结果在 R5
            EmitMov(REG_TEMP, REG_ACCUM);   // 将右操作数的值 (当前在R5) 移动到 REG_TEMP (R6)
            EmitPop(REG_ACCUM);             // 从栈中弹出之前保存的左操作数到 R5

            // 获取实际的操作符类型 (MUL 或 DIV)
            var opTokenNode = context.mul_div_op().GetChild(0) as ITerminalNode;
            if (opTokenNode.Symbol.Type == TinyLanguageParser.MUL)
            {
                EmitOpRegs(OpCode.MUL, REG_ACCUM, REG_ACCUM, REG_TEMP); // R5 = R5 * R6
            }
            else if (opTokenNode.Symbol.Type == TinyLanguageParser.DIV)
            {
                EmitOpRegs(OpCode.DIV, REG_ACCUM, REG_ACCUM, REG_TEMP); // R5 = R5 / R6
            }
            return null; // 最终结果已存放在 R5 (REG_ACCUM)
        }

        /// <summary>
        /// 访问加减法表达式节点 (AddSubExpr)。
        /// </summary>
        public override object VisitAddSubExpr(TinyLanguageParser.AddSubExprContext context)
        {
            Visit(context.expression(0));   // 计算左操作数 -> R5
            EmitPush(REG_ACCUM);            // PUSH R5 (左)
            Visit(context.expression(1));   // 计算右操作数 -> R5
            EmitMov(REG_TEMP, REG_ACCUM);   // R6 = R5 (右)
            EmitPop(REG_ACCUM);             // POP R5 (左)

            var opTokenNode = context.add_sub_op().GetChild(0) as ITerminalNode;
            if (opTokenNode.Symbol.Type == TinyLanguageParser.ADD)
            {
                EmitOpRegs(OpCode.ADD, REG_ACCUM, REG_ACCUM, REG_TEMP); // R5 = R5 + R6
            }
            else if (opTokenNode.Symbol.Type == TinyLanguageParser.SUB)
            {
                EmitOpRegs(OpCode.SUB, REG_ACCUM, REG_ACCUM, REG_TEMP); // R5 = R5 - R6
            }
            return null; // 结果在 R5
        }

        /// <summary>
        /// 访问表达式中的项节点 (TermExpr)。只是简单地委托给具体的项类型。
        /// </summary>
        public override object VisitTermExpr(TinyLanguageParser.TermExprContext context)
        {
            return Visit(context.term());
        }

        /// <summary>
        /// 访问整数常量项节点 (IntegerTerm)。
        /// </summary>
        public override object VisitIntegerTerm(TinyLanguageParser.IntegerTermContext context)
        {
            short value = short.Parse(context.INTEGER().GetText()); // 将文本转换为 short 类型整数
            EmitConst(REG_ACCUM, value); // 将常量加载到 REG_ACCUM (R5)
            return null;
        }

        /// <summary>
        /// 访问标识符项节点 (IdentifierTerm)，即变量。
        /// </summary>
        public override object VisitIdentifierTerm(TinyLanguageParser.IdentifierTermContext context)
        {
            string varName = context.IDENTIFIER().GetText();
            if (!symbolTable.TryGetValue(varName, out SymbolTableEntry entry))
            {
                throw new Exception($"语义错误: 变量 '{varName}' 未声明，不能在表达式中使用。");
            }
            EmitLoad(REG_ACCUM, entry.Address); // 从内存加载变量值到 REG_ACCUM (R5)
            return null;
        }

        /// <summary>
        /// 访问括号表达式项节点 (ParenExpr)。
        /// </summary>
        public override object VisitParenExpr(TinyLanguageParser.ParenExprContext context)
        {
            return Visit(context.expression()); // 计算括号内部表达式的值，结果会存入 REG_ACCUM (R5)
        }

        // --- 条件求值 ---
        // 约定：条件求值后，其布尔结果 (0 代表 false, 1 代表 true) 存放在 REG_COND_RESULT (R4) 寄存器中。
        public override object VisitCondition(TinyLanguageParser.ConditionContext context)
        {
            // 1. 计算左表达式，结果在 R5 (REG_ACCUM)
            Visit(context.expression(0));
            EmitPush(REG_ACCUM); // 保存左表达式结果到栈

            // 2. 计算右表达式，结果在 R5 (REG_ACCUM)
            Visit(context.expression(1));
            EmitMov(REG_TEMP, REG_ACCUM); // 将右表达式结果 (当前在R5) 移到 R6 (REG_TEMP)

            // 3. 从栈恢复左表达式结果到 R5
            EmitPop(REG_ACCUM); // R5 = 左表达式结果

            // 4. 计算 (左表达式 - 右表达式)，结果存入 R3 (REG_CMP_TEMP)
            EmitOpRegs(OpCode.SUB, REG_CMP_TEMP, REG_ACCUM, REG_TEMP); // R3 = R5 - R6

            // 5. 默认将条件结果 R4 (REG_COND_RESULT) 设置为 0 (false)
            EmitConst(REG_COND_RESULT, 0);

            List<int> setTrueLabel = CreateLabelPlaceholderList();    // 跳转到此标签处，将 R4 设置为 1 (true)
            List<int> endConditionLabel = CreateLabelPlaceholderList(); // 条件评估结束的标签

            var opToken = context.relational_op().GetChild(0) as ITerminalNode; // 获取关系操作符
            OpCode vmConditionalJumpOp = OpCode.HLT; // 用于 EQ, NEQ, LT, GT 的VM条件跳转指令 (默认为HLT，应被覆盖)

            // 根据TL的关系操作符选择VM的条件跳转指令，并生成相应的跳转逻辑
            switch (opToken.Symbol.Type)
            {
                case TinyLanguageParser.EQ:  // 如果 (R5 == R6)，即 (R3 == 0)，则条件为真
                    vmConditionalJumpOp = OpCode.JZ;
                    EmitCondJmp(vmConditionalJumpOp, REG_CMP_TEMP, setTrueLabel); // JZ R3, setTrueLabel
                    EmitJmp(endConditionLabel); // 如果不等于0，则R4保持0，跳到结束
                    break;
                case TinyLanguageParser.NEQ: // 如果 (R5 != R6)，即 (R3 != 0)，则条件为真
                    vmConditionalJumpOp = OpCode.JNZ;
                    EmitCondJmp(vmConditionalJumpOp, REG_CMP_TEMP, setTrueLabel); // JNZ R3, setTrueLabel
                    EmitJmp(endConditionLabel);
                    break;
                case TinyLanguageParser.LT:  // 如果 (R5 < R6)，即 (R3 < 0)，则条件为真
                    vmConditionalJumpOp = OpCode.JL;
                    EmitCondJmp(vmConditionalJumpOp, REG_CMP_TEMP, setTrueLabel); // JL R3, setTrueLabel
                    EmitJmp(endConditionLabel);
                    break;
                case TinyLanguageParser.GT:  // 如果 (R5 > R6)，即 (R3 > 0)，则条件为真
                    vmConditionalJumpOp = OpCode.JG;
                    EmitCondJmp(vmConditionalJumpOp, REG_CMP_TEMP, setTrueLabel); // JG R3, setTrueLabel
                    EmitJmp(endConditionLabel);
                    break;
                case TinyLanguageParser.LTE: // 如果 (R5 <= R6)，即 NOT (R5 > R6)，即 NOT (R3 > 0)
                                   // 如果 R3 > 0 (JG R3)，则条件为假 (R4保持0)，直接跳到结束。
                                   // 否则 (R3 <= 0)，条件为真，需要跳转到设置R4=1的地方。
                    EmitCondJmp(OpCode.JG, REG_CMP_TEMP, endConditionLabel); // 如果 R5 > R6, 跳转到结束 (R4=0)
                    EmitJmp(setTrueLabel); // 否则 (R5 <= R6), 跳转去设置 R4=1
                    break;
                case TinyLanguageParser.GTE: // 如果 (R5 >= R6)，即 NOT (R5 < R6)，即 NOT (R3 < 0)
                                   // 如果 R3 < 0 (JL R3)，则条件为假 (R4保持0)，直接跳到结束。
                                   // 否则 (R3 >= 0)，条件为真，需要跳转到设置R4=1的地方。
                    EmitCondJmp(OpCode.JL, REG_CMP_TEMP, endConditionLabel); // 如果 R5 < R6, 跳转到结束 (R4=0)
                    EmitJmp(setTrueLabel); // 否则 (R5 >= R6), 跳转去设置 R4=1
                    break;
                default:
                    throw new Exception($"未知的关系操作符: {opToken.GetText()}");
            }

            PlaceLabel(setTrueLabel);    // 定义 setTrueLabel 的位置
            EmitConst(REG_COND_RESULT, 1); // 在此路径上，将 R4 设置为 1 (true)
                                           // 此处可以隐式地“落到”endConditionLabel，或者再加一个JMP endConditionLabel确保路径清晰
                                           // 为了结构统一，可以考虑在setTrueLabel后也JMP到endConditionLabel，但如果endConditionLabel紧随其后则非必须。
                                           // 当前的实现是，如果跳到setTrueLabel，执行完CONST R4,1后，会顺序执行到PlaceLabel(endConditionLabel)。

            PlaceLabel(endConditionLabel); // 定义 endConditionLabel 的位置
            // 至此，REG_COND_RESULT (R4) 中已保存了条件的布尔结果 (0 或 1)
            return null;
        }


        // --- 控制流语句 ---

        /// <summary>
        /// 访问 IF 语句节点 (if_statement)。
        /// </summary>
        public override object VisitIf_statement(TinyLanguageParser.If_statementContext context)
        {
            Visit(context.condition()); // 评估条件，结果 (0或1) 存入 REG_COND_RESULT (R4)

            List<int> elseLabel = CreateLabelPlaceholderList();   // 'else' 块或 'if' 结束点（若无else）的标签占位符
            List<int> endIfLabel = CreateLabelPlaceholderList(); // 'if' 语句完全结束点的标签占位符 (主要用于 'then' 块跳过 'else' 块)

            // 如果条件为假 (R4 == 0), 则跳转到 elseLabel
            EmitCondJmp(OpCode.JZ, REG_COND_RESULT, elseLabel);

            // Then 块: 如果条件为真，则执行此处的语句
            Visit(context.statement_list(0)); // 访问 'then' 块的语句列表

            var elseClauseNode = context.K_ELSE(); // 检查是否存在 ELSE 子句
            if (elseClauseNode != null)
            {
                // 如果存在 ELSE 子句，那么 'then' 块执行完毕后需要无条件跳转到整个 IF 语句的末尾
                EmitJmp(endIfLabel);
            }

            PlaceLabel(elseLabel); // 定义 elseLabel 的位置
                                   // 如果没有 ELSE 子句，这个标签实际上就是 IF 语句的结束点

            if (elseClauseNode != null)
            {
                Visit(context.statement_list(1)); // 访问 'else' 块的语句列表
            }

            PlaceLabel(endIfLabel); // 定义整个 IF 语句结束的标签 (即使没有 ELSE，这个标签也会被定义，
                                    // 只是可能与 elseLabel 指向同一位置)
            return null;
        }

        /// <summary>
        /// 访问 WHILE 语句节点 (while_statement)。
        /// </summary>
        public override object VisitWhile_statement(TinyLanguageParser.While_statementContext context)
        {
            List<int> loopEndLabel = CreateLabelPlaceholderList(); // 循环结束的标签占位符

            ushort loopStartAddress = (ushort)bytecode.Count; // 记录循环开始的实际地址 (即条件判断的开始)

            Visit(context.condition()); // 评估条件，结果 (0或1) 存入 REG_COND_RESULT (R4)

            // 如果条件为假 (R4 == 0), 则跳转到循环结束 (loopEndLabel)
            EmitCondJmp(OpCode.JZ, REG_COND_RESULT, loopEndLabel);

            // 循环体: 如果条件为真，则执行此处的语句
            Visit(context.statement_list()); // 访问循环体的语句列表
                                             // 注意: TL.g4 中 while_statement只有一个 statement_list,
                                             // ANTLR 生成的 context.statement_list() 会直接返回该节点
                                             // 而不是像 if 中那样返回一个列表，所以不需要索引 [0]

            // 无条件跳转回循环开始 (重新进行条件判断)
            EmitOpCode(OpCode.JMP);
            EmitUShort(loopStartAddress); // 直接跳转到已记录的循环起始地址

            PlaceLabel(loopEndLabel); // 定义循环结束标签的位置
            return null;
        }
    }
}