diff --git a/src/LLVMIRGenerator.cpp b/src/LLVMIRGenerator.cpp
index b6494d9..32c4a23 100644
--- a/src/LLVMIRGenerator.cpp
+++ b/src/LLVMIRGenerator.cpp
@@ -217,7 +217,7 @@ std::any LLVMIRGenerator::visitBlockStmt(SysYParser::BlockStmtContext* ctx) {
     }
     return nullptr;
 }
-std::any visitAssignStmt(SysYParser::AssignStmtContext *ctx)
+std::any LLVMIRGenerator::visitAssignStmt(SysYParser::AssignStmtContext *ctx)
 {
     std::string lhsAlloca = std::any_cast<std::string>(ctx->lValue()->accept(this));
     std::string lhsType = symbolTable[ctx->lValue()->Ident()->getText()].second;
@@ -240,7 +240,7 @@ std::any visitAssignStmt(SysYParser::AssignStmtContext *ctx)
     return nullptr;
 }
 
-std::any visitIfStmt(SysYParser::IfStmtContext *ctx)
+std::any LLVMIRGenerator::visitIfStmt(SysYParser::IfStmtContext *ctx)
 {
     std::string cond = std::any_cast<std::string>(ctx->cond()->accept(this));
     std::string trueLabel = "if.then." + std::to_string(tempCounter);
@@ -263,7 +263,7 @@ std::any visitIfStmt(SysYParser::IfStmtContext *ctx)
     return nullptr;
 }
 
-std::any visitWhileStmt(SysYParser::WhileStmtContext *ctx)
+std::any LLVMIRGenerator::visitWhileStmt(SysYParser::WhileStmtContext *ctx)
 {
     std::string loop_cond = "while.cond." + std::to_string(tempCounter);
     std::string loop_body = "while.body." + std::to_string(tempCounter);
@@ -276,7 +276,7 @@ std::any visitWhileStmt(SysYParser::WhileStmtContext *ctx)
     std::string cond = std::any_cast<std::string>(ctx->cond()->accept(this));
     irStream << "  br i1 " << cond << ", label %" << loop_body << ", label %" << loop_end << "\n";
     irStream << loop_body << ":\n";
-    ctx->stmt(0)->accept(this);
+    ctx->stmt()->accept(this);
     irStream << "  br label %" << loop_cond << "\n";
     irStream << loop_end << ":\n";
     
@@ -284,7 +284,7 @@ std::any visitWhileStmt(SysYParser::WhileStmtContext *ctx)
     return nullptr;
 }
 
-std::any visitBreakStmt(SysYParser::BreakStmtContext *ctx)
+std::any LLVMIRGenerator::visitBreakStmt(SysYParser::BreakStmtContext *ctx)
 {
     if (loopStack.empty()) {
         throw std::runtime_error("Break statement outside of a loop.");
@@ -293,7 +293,7 @@ std::any visitBreakStmt(SysYParser::BreakStmtContext *ctx)
     return nullptr;
 }
 
-std::any visitContinueStmt(SysYParser::ContinueStmtContext *ctx)
+std::any LLVMIRGenerator::visitContinueStmt(SysYParser::ContinueStmtContext *ctx)
 {
     if (loopStack.empty()) {
         throw std::runtime_error("Continue statement outside of a loop.");
@@ -302,6 +302,18 @@ std::any visitContinueStmt(SysYParser::ContinueStmtContext *ctx)
     return nullptr;
 }
 
+std::any LLVMIRGenerator::visitReturnStmt(SysYParser::ReturnStmtContext *ctx)
+{
+    hasReturn = true;
+    if (ctx->exp()) {
+        std::string value = std::any_cast<std::string>(ctx->exp()->accept(this));
+        irStream << "  ret " << currentReturnType << " " << value << "\n";
+    } else {
+        irStream << "  ret void\n";
+    }
+    return nullptr;
+}
+
 // std::any LLVMIRGenerator::visitStmt(SysYParser::StmtContext* ctx) {
 //     if (ctx->lValue() && ctx->ASSIGN()) {
 //         std::string lhsAlloca = std::any_cast<std::string>(ctx->lValue()->accept(this));
@@ -385,23 +397,34 @@ std::any visitContinueStmt(SysYParser::ContinueStmtContext *ctx)
 
 std::any LLVMIRGenerator::visitLValue(SysYParser::LValueContext* ctx) {
     std::string varName = ctx->Ident()->getText();
-    return symbolTable[varName].first;
+    std::string allocaPtr = symbolTable[varName].first;
+    std::string type = symbolTable[varName].second;
+    std::string temp = getNextTemp();
+    irStream << "  " << temp << " = load " << type << ", " << type << "* " << allocaPtr << ", align 4\n";
+    tmpTable[temp] = type;
+    return temp;
+    // std::string varName = ctx->Ident()->getText();
+    // return symbolTable[varName].first;
 }
 
-std::any LLVMIRGenerator::visitPrimaryExp(SysYParser::PrimaryExpContext* ctx) {
-    if (ctx->lValue()) {
-        std::string allocaPtr = std::any_cast<std::string>(ctx->lValue()->accept(this));
-        std::string varName = ctx->lValue()->Ident()->getText();
-        std::string type = symbolTable[varName].second;
-        std::string temp = getNextTemp();
-        irStream << "  " << temp << " = load " << type << ", " << type << "* " << allocaPtr << ", align 4\n";
-        tmpTable[temp] = type;
-        return temp;
-    } else if (ctx->exp()) {
-        return ctx->exp()->accept(this);
-    } else {
-        return ctx->number()->accept(this);
-    }
+// std::any LLVMIRGenerator::visitPrimaryExp(SysYParser::PrimaryExpContext* ctx) {
+//     if (ctx->lValue()) {
+//         std::string allocaPtr = std::any_cast<std::string>(ctx->lValue()->accept(this));
+//         std::string varName = ctx->lValue()->Ident()->getText();
+//         std::string type = symbolTable[varName].second;
+//         std::string temp = getNextTemp();
+//         irStream << "  " << temp << " = load " << type << ", " << type << "* " << allocaPtr << ", align 4\n";
+//         tmpTable[temp] = type;
+//         return temp;
+//     } else if (ctx->exp()) {
+//         return ctx->exp()->accept(this);
+//     } else {
+//         return ctx->number()->accept(this);
+//     }
+// }
+
+std::any LLVMIRGenerator::visitPrimExp(SysYParser::PrimaryExpContext* ctx) {
+  return visitChildren(ctx);
 }
 
 std::any LLVMIRGenerator::visitNumber(SysYParser::NumberContext* ctx) {
@@ -413,7 +436,7 @@ std::any LLVMIRGenerator::visitNumber(SysYParser::NumberContext* ctx) {
     return "";
 }
 
-std::any LLVMIRGenerator::visitUnaryExp(SysYParser::UnaryExpContext* ctx) {
+std::any LLVMIRGenerator::visitUnExp(SysYParser::UnExpContext* ctx) {
     if (ctx->unaryOp()) {
         std::string operand = std::any_cast<std::string>(ctx->unaryExp()->accept(this));
         std::string op = ctx->unaryOp()->getText();
@@ -426,25 +449,27 @@ std::any LLVMIRGenerator::visitUnaryExp(SysYParser::UnaryExpContext* ctx) {
             irStream << "  " << temp << " = xor " << type << " " << operand << ", 1\n";
         }
         return temp;
-    } else if (ctx->Ident()) {
-        std::string funcName = ctx->Ident()->getText();
-        std::vector<std::string> args;
-        if (ctx->funcRParams()) {
-            for (auto argCtx : ctx->funcRParams()->exp()) {
-                args.push_back(std::any_cast<std::string>(argCtx->accept(this)));
-            }
-        }
-        std::string temp = getNextTemp();
-        std::string argList = "";
-        for (size_t i = 0; i < args.size(); ++i) {
-            if (i > 0) argList += ", ";
-            argList +=tmpTable[args[i]] + " noundef " + args[i];
-        }
-        irStream << "  " << temp << " = call " << currentReturnType << " @" << funcName << "(" << argList << ")\n";
-        tmpTable[temp] = currentReturnType;
-        return temp;
     }
-    return ctx->primaryExp()->accept(this);
+    return ctx->unaryExp()->accept(this);
+}
+std::any LLVMIRGenerator::visitCall(SysYParser::CallContext *ctx)
+{
+    std::string funcName = ctx->Ident()->getText();
+    std::vector<std::string> args;
+    if (ctx->funcRParams()) {
+        for (auto argCtx : ctx->funcRParams()->exp()) {
+            args.push_back(std::any_cast<std::string>(argCtx->accept(this)));
+        }
+    }
+    std::string temp = getNextTemp();
+    std::string argList = "";
+    for (size_t i = 0; i < args.size(); ++i) {
+        if (i > 0) argList += ", ";
+        argList +=tmpTable[args[i]] + " noundef " + args[i];
+    }
+    irStream << "  " << temp << " = call " << currentReturnType << " @" << funcName << "(" << argList << ")\n";
+    tmpTable[temp] = currentReturnType;
+    return temp;
 }
 
 std::any LLVMIRGenerator::visitMulExp(SysYParser::MulExpContext* ctx) {
diff --git a/src/LLVMIRGenerator.h b/src/LLVMIRGenerator.h
index f5b9430..a02419f 100644
--- a/src/LLVMIRGenerator.h
+++ b/src/LLVMIRGenerator.h
@@ -41,11 +41,12 @@ private:
     std::any visitVarDef(SysYParser::VarDefContext* ctx);
     std::any visitFuncDef(SysYParser::FuncDefContext* ctx);
     std::any visitBlockStmt(SysYParser::BlockStmtContext* ctx);
-    std::any visitStmt(SysYParser::StmtContext* ctx);
+    // std::any visitStmt(SysYParser::StmtContext* ctx);
     std::any visitLValue(SysYParser::LValueContext* ctx);
-    std::any visitPrimaryExp(SysYParser::PrimaryExpContext* ctx);
+    std::any visitPrimExp(SysYParser::PrimaryExpContext* ctx);
     std::any visitNumber(SysYParser::NumberContext* ctx);
-    std::any visitUnaryExp(SysYParser::UnaryExpContext* ctx);
+    std::any visitCall(SysYParser::CallContext *ctx);
+    std::any visitUnExp(SysYParser::UnExpContext* ctx);
     std::any visitMulExp(SysYParser::MulExpContext* ctx);
     std::any visitAddExp(SysYParser::AddExpContext* ctx);
     std::any visitRelExp(SysYParser::RelExpContext* ctx);
diff --git a/src/sysyc.cpp b/src/sysyc.cpp
index 6ec6eb7..b509fab 100644
--- a/src/sysyc.cpp
+++ b/src/sysyc.cpp
@@ -78,24 +78,26 @@ int main(int argc, char **argv) {
   }
 
   // visit AST to generate IR
-  SysYIRGenerator generator;
-  generator.visitCompUnit(moduleAST);
-  auto moduleIR = generator.get();
+
   if (argStopAfter == "ir") {
+    SysYIRGenerator generator;
+    generator.visitCompUnit(moduleAST);
+    auto moduleIR = generator.get();
     moduleIR->print(cout);
     return EXIT_SUCCESS;
   } else if (argStopAfter == "llvmir") {
     LLVMIRGenerator llvmirGenerator;
+    llvmirGenerator.generateIR(moduleAST); // 使用公共接口生成 IR
     cout << llvmirGenerator.getIR();
     return EXIT_SUCCESS;
   }
 
-  // generate assembly
-  CodeGen codegen(moduleIR);
-  string asmCode = codegen.code_gen();
-  cout << asmCode << endl;
-  if (argStopAfter == "asm")
-    return EXIT_SUCCESS;
+  // // generate assembly
+  // CodeGen codegen(moduleIR);
+  // string asmCode = codegen.code_gen();
+  // cout << asmCode << endl;
+  // if (argStopAfter == "asm")
+  //   return EXIT_SUCCESS;
 
   return EXIT_SUCCESS;
 }
\ No newline at end of file