Added STRING

4 files changed, 319 insertions, 1 deletions

diff --git a/lib/CodeGen/CodeGen.cc b/lib/CodeGen/CodeGen.cc
index 95e3768..ee7ab04 100644
--- a/lib/CodeGen/CodeGen.cc
+++ b/lib/CodeGen/CodeGen.cc
@@ -24,6 +24,7 @@ CodeGen::CodeGen() {
     RealHelper = std::make_unique<RealHandler>(*TheContext, *Builder, *TheModule, NamedValues);
     BoolHandler = std::make_unique<BooleanHandler>(*TheContext, *Builder, *TheModule, NamedValues);
     ArithHandler = std::make_unique<ArithmeticHandler>(*TheContext, *Builder);
+    StrHandler = std::make_unique<StringHandler>(*TheContext, *Builder, *TheModule, NamedValues);
 
     SetupExternalFunctions();
 }
@@ -43,6 +44,7 @@ void CodeGen::SetupExternalFunctions() {
 
     ScanfFormatStr = Builder->CreateGlobalStringPtr("%lld", "fmt_in", 0, TheModule.get());
     ScanfFloatFormatStr = Builder->CreateGlobalStringPtr("%lf", "fmt_in_flt", 0, TheModule.get());
+    ScanfStringFormatStr = Builder->CreateGlobalStringPtr("%s", "fmt_in_str", 0, TheModule.get()); // NEW
 
     TrueStr = Builder->CreateGlobalStringPtr("TRUE", "str_true", 0, TheModule.get());
     FalseStr = Builder->CreateGlobalStringPtr("FALSE", "str_false", 0, TheModule.get());
@@ -89,6 +91,10 @@ Value *CodeGen::emitExpr(ExprAST *Expr) {
     if (auto *Bool = dynamic_cast<BooleanExprAST*>(Expr)) {
         return BoolHandler->createLiteral(Bool->getVal());
     }
+
+    if (auto *Str = dynamic_cast<StringExprAST*>(Expr)) {
+        return StrHandler->createLiteral(Str->getVal());
+    }
     
     if (auto *Var = dynamic_cast<VariableExprAST*>(Expr)) {
         Value *A = NamedValues[Var->getName()];
@@ -129,6 +135,31 @@ Value *CodeGen::emitExpr(ExprAST *Expr) {
     }
 
     if (auto *Call = dynamic_cast<CallExprAST*>(Expr)) {
+        std::string Name = Call->getCallee();
+        if (Name == "LENGTH") {
+            if (Call->getArgs().size() != 1) { fprintf(stderr, "LENGTH expects 1 arg\n"); return nullptr; }
+            return StrHandler->emitLength(emitExpr(Call->getArgs()[0].get()));
+        }
+        if (Name == "MID") {
+            if (Call->getArgs().size() != 3) { fprintf(stderr, "MID expects 3 args\n"); return nullptr; }
+            return StrHandler->emitMid(emitExpr(Call->getArgs()[0].get()), 
+                                       emitExpr(Call->getArgs()[1].get()), 
+                                       emitExpr(Call->getArgs()[2].get()));
+        }
+        if (Name == "RIGHT") {
+            if (Call->getArgs().size() != 2) { fprintf(stderr, "RIGHT expects 2 args\n"); return nullptr; }
+            return StrHandler->emitRight(emitExpr(Call->getArgs()[0].get()), 
+                                         emitExpr(Call->getArgs()[1].get()));
+        }
+        if (Name == "LCASE") {
+            if (Call->getArgs().size() != 1) { fprintf(stderr, "LCASE expects 1 arg\n"); return nullptr; }
+            return StrHandler->emitLCase(emitExpr(Call->getArgs()[0].get()));
+        }
+        if (Name == "UCASE") {
+            if (Call->getArgs().size() != 1) { fprintf(stderr, "UCASE expects 1 arg\n"); return nullptr; }
+            return StrHandler->emitUCase(emitExpr(Call->getArgs()[0].get()));
+        }
+
         Function *CalleeF = TheModule->getFunction(Call->getCallee());
         std::vector<Value*> Args;
         
@@ -382,6 +413,7 @@ void CodeGen::emitStmt(StmtAST *Stmt) {
         if (Decl->getType() == "INTEGER") IntHandler->emitDeclare(Decl->getName());
         else if (Decl->getType() == "REAL") RealHelper->emitDeclare(Decl->getName());
         else if (Decl->getType() == "BOOLEAN") BoolHandler->emitDeclare(Decl->getName());
+        else if (Decl->getType() == "STRING") StrHandler->emitDeclare(Decl->getName());
         else {
             fprintf(stderr, "Unknown type %s\n", Decl->getType().c_str());
         }
@@ -429,6 +461,22 @@ void CodeGen::emitStmt(StmtAST *Stmt) {
                  Value *BoolVal = Builder->CreateICmpNE(Val, ConstantInt::get(*TheContext, APInt(64, 0)), "bool_cast");
                  
                  Builder->CreateStore(BoolVal, Alloca);
+            } else if (AI->getAllocatedType()->isPointerTy()) {
+                 Value *BufSize = ConstantInt::get(*TheContext, APInt(64, 256));
+
+                 Args.push_back(ScanfStringFormatStr);
+
+                 AllocaInst *InputBuf = Builder->CreateAlloca(Type::getInt8Ty(*TheContext), ConstantInt::get(*TheContext, APInt(64, 1024)), "input_buf");
+                 Args.push_back(InputBuf);
+                 Builder->CreateCall(ScanfFunc, Args);
+
+                 Function *MallocF = TheModule->getFunction("malloc");
+                 if (MallocF) {
+                     Value *Mem = Builder->CreateCall(MallocF, ConstantInt::get(*TheContext, APInt(64, 1024)));
+                     Args[1] = Mem;
+                     Builder->CreateCall(ScanfFunc, Args);
+                     Builder->CreateStore(Mem, Alloca);
+                 }
             } else {
                  Args.push_back(ScanfFormatStr);
                  Args.push_back(Alloca); 
@@ -450,6 +498,9 @@ void CodeGen::emitStmt(StmtAST *Stmt) {
                      Args.push_back(PrintfStringFormatStr);
                      Value *StrPtr = Builder->CreateSelect(Val, TrueStr, FalseStr);
                      Args.push_back(StrPtr);
+                } else if (Val->getType()->isPointerTy()) {
+                     Args.push_back(PrintfStringFormatStr);
+                     Args.push_back(Val);
                 } else {
                      Args.push_back(PrintfFormatStr);
                      Args.push_back(Val);
diff --git a/lib/CodeGen/StringHandler.cc b/lib/CodeGen/StringHandler.cc
new file mode 100644
index 0000000..fda9ea5
--- /dev/null
+++ b/lib/CodeGen/StringHandler.cc
@@ -0,0 +1,209 @@
+#include "cps/StringHandler.h"
+#include "llvm/IR/Type.h"
+#include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/Constants.h"
+
+using namespace llvm;
+using namespace cps;
+
+StringHandler::StringHandler(LLVMContext &Ctx, IRBuilder<> &B, llvm::Module &M, 
+                             std::map<std::string, Value*> &NV)
+    : Context(Ctx), Builder(B), Module(M), NamedValues(NV) {
+    setupExternalFunctions();
+}
+
+void StringHandler::setupExternalFunctions() {
+    FunctionType *MallocType = FunctionType::get(PointerType::getUnqual(Context), {Type::getInt64Ty(Context)}, false);
+    MallocFunc = Module.getOrInsertFunction("malloc", MallocType);
+
+    FunctionType *FreeType = FunctionType::get(Type::getVoidTy(Context), {PointerType::getUnqual(Context)}, false);
+    FreeFunc = Module.getOrInsertFunction("free", FreeType);
+
+    FunctionType *StrLenType = FunctionType::get(Type::getInt64Ty(Context), {PointerType::getUnqual(Context)}, false);
+    StrLenFunc = Module.getOrInsertFunction("strlen", StrLenType);
+
+    std::vector<Type*> MemCpyArgs = {PointerType::getUnqual(Context), PointerType::getUnqual(Context), Type::getInt64Ty(Context)};
+    FunctionType *MemCpyType = FunctionType::get(PointerType::getUnqual(Context), MemCpyArgs, false);
+    MemCpyFunc = Module.getOrInsertFunction("memcpy", MemCpyType);
+
+    FunctionType *CharCaseType = FunctionType::get(Type::getInt32Ty(Context), {Type::getInt32Ty(Context)}, false);
+    ToUpperFunc = Module.getOrInsertFunction("toupper", CharCaseType);
+    ToLowerFunc = Module.getOrInsertFunction("tolower", CharCaseType);
+}
+
+void StringHandler::emitDeclare(const std::string &Name) {
+    Function *TheFunction = Builder.GetInsertBlock()->getParent();
+    IRBuilder<> TmpB(&TheFunction->getEntryBlock(), TheFunction->getEntryBlock().begin());
+    AllocaInst *Alloca = TmpB.CreateAlloca(PointerType::getUnqual(Context), nullptr, Name);
+
+    Value *EmptyStr = createLiteral("");
+    Builder.CreateStore(EmptyStr, Alloca);
+    NamedValues[Name] = Alloca;
+}
+
+Value *StringHandler::createLiteral(const std::string &Val) {
+    return Builder.CreateGlobalStringPtr(Val);
+}
+
+Value *StringHandler::emitLength(Value *Str) {
+    if (!Str) return nullptr;
+    return Builder.CreateCall(StrLenFunc, Str, "len");
+}
+
+Value *StringHandler::emitMid(Value *Str, Value *Start, Value *Len) {
+    
+    Value *FullLen = emitLength(Str);
+
+    Value *One = ConstantInt::get(Context, APInt(64, 1));
+    Value *Zero = ConstantInt::get(Context, APInt(64, 0));
+    Value *StartZeroBased = Builder.CreateSub(Start, One, "start_idx");
+
+    Value *IsNeg = Builder.CreateICmpSLT(StartZeroBased, Zero);
+    StartZeroBased = Builder.CreateSelect(IsNeg, Zero, StartZeroBased);
+
+    Value *IsTooBig = Builder.CreateICmpSGT(StartZeroBased, FullLen);
+    StartZeroBased = Builder.CreateSelect(IsTooBig, FullLen, StartZeroBased);
+
+    Value *RemLen = Builder.CreateSub(FullLen, StartZeroBased);
+
+    Value *IsLenTooBig = Builder.CreateICmpSGT(Len, RemLen);
+    Value *ActualLen = Builder.CreateSelect(IsLenTooBig, RemLen, Len);
+
+    Value *IsLenNeg = Builder.CreateICmpSLT(ActualLen, Zero);
+    ActualLen = Builder.CreateSelect(IsLenNeg, Zero, ActualLen);
+
+    Value *AllocSize = Builder.CreateAdd(ActualLen, One);
+    Value *NewStrMem = Builder.CreateCall(MallocFunc, AllocSize, "mid_str_mem");
+
+    Value *SrcPtr = Builder.CreateInBoundsGEP(Type::getInt8Ty(Context), Str, StartZeroBased);
+
+    std::vector<Value*> Args = {NewStrMem, SrcPtr, ActualLen};
+    Builder.CreateCall(MemCpyFunc, Args);
+
+    Value *NullTermPtr = Builder.CreateInBoundsGEP(Type::getInt8Ty(Context), NewStrMem, ActualLen);
+    Builder.CreateStore(ConstantInt::get(Type::getInt8Ty(Context), 0), NullTermPtr);
+    
+    return NewStrMem;
+}
+
+Value *StringHandler::emitRight(Value *Str, Value *Len) {
+    Value *FullLen = emitLength(Str);
+
+    Value *StartIdx = Builder.CreateSub(FullLen, Len, "right_start");
+
+    Value *Zero = ConstantInt::get(Context, APInt(64, 0));
+    Value *IsNeg = Builder.CreateICmpSLT(StartIdx, Zero);
+    StartIdx = Builder.CreateSelect(IsNeg, Zero, StartIdx);
+
+    Value *ActualLen = Builder.CreateSub(FullLen, StartIdx);
+
+    Value *One = ConstantInt::get(Context, APInt(64, 1));
+    Value *AllocSize = Builder.CreateAdd(ActualLen, One);
+    Value *NewStrMem = Builder.CreateCall(MallocFunc, AllocSize, "right_str_mem");
+
+    Value *SrcPtr = Builder.CreateInBoundsGEP(Type::getInt8Ty(Context), Str, StartIdx);
+
+    std::vector<Value*> Args = {NewStrMem, SrcPtr, ActualLen};
+    Builder.CreateCall(MemCpyFunc, Args);
+
+    Value *NullTermPtr = Builder.CreateInBoundsGEP(Type::getInt8Ty(Context), NewStrMem, ActualLen);
+    Builder.CreateStore(ConstantInt::get(Type::getInt8Ty(Context), 0), NullTermPtr);
+    
+    return NewStrMem;
+}
+
+Value *StringHandler::emitLCase(Value *Str) {
+    Value *Len = emitLength(Str);
+    Value *One = ConstantInt::get(Context, APInt(64, 1));
+    Value *AllocSize = Builder.CreateAdd(Len, One);
+    Value *NewStr = Builder.CreateCall(MallocFunc, AllocSize, "lcase_str");
+
+    Function *TheFunction = Builder.GetInsertBlock()->getParent();
+    BasicBlock *LoopBB = BasicBlock::Create(Context, "loop", TheFunction);
+    BasicBlock *AfterBB = BasicBlock::Create(Context, "afterloop", TheFunction);
+
+    IRBuilder<> TmpB(&TheFunction->getEntryBlock(), TheFunction->getEntryBlock().begin());
+    AllocaInst *IdxVar = TmpB.CreateAlloca(Type::getInt64Ty(Context), nullptr, "idx");
+    Builder.CreateStore(ConstantInt::get(Context, APInt(64, 0)), IdxVar);
+    
+    Builder.CreateBr(LoopBB);
+    Builder.SetInsertPoint(LoopBB);
+    
+    Value *CurIdx = Builder.CreateLoad(Type::getInt64Ty(Context), IdxVar);
+    Value *Cond = Builder.CreateICmpSLT(CurIdx, Len);
+    
+    BasicBlock *BodyBB = BasicBlock::Create(Context, "body", TheFunction);
+    Builder.CreateCondBr(Cond, BodyBB, AfterBB);
+    
+    Builder.SetInsertPoint(BodyBB);
+
+    Value *SrcPtr = Builder.CreateInBoundsGEP(Type::getInt8Ty(Context), Str, CurIdx);
+    Value *CharVal = Builder.CreateLoad(Type::getInt8Ty(Context), SrcPtr);
+
+    Value *ExtChar = Builder.CreateSExt(CharVal, Type::getInt32Ty(Context));
+    Value *LowerChar = Builder.CreateCall(ToLowerFunc, ExtChar);
+    Value *TruncChar = Builder.CreateTrunc(LowerChar, Type::getInt8Ty(Context));
+
+    Value *DestPtr = Builder.CreateInBoundsGEP(Type::getInt8Ty(Context), NewStr, CurIdx);
+    Builder.CreateStore(TruncChar, DestPtr);
+
+    Value *NextIdx = Builder.CreateAdd(CurIdx, One);
+    Builder.CreateStore(NextIdx, IdxVar);
+    Builder.CreateBr(LoopBB);
+    
+    Builder.SetInsertPoint(AfterBB);
+
+    Value *NullPtr = Builder.CreateInBoundsGEP(Type::getInt8Ty(Context), NewStr, Len);
+    Builder.CreateStore(ConstantInt::get(Type::getInt8Ty(Context), 0), NullPtr);
+    
+    return NewStr;
+}
+
+Value *StringHandler::emitUCase(Value *Str) {
+    Value *Len = emitLength(Str);
+    Value *One = ConstantInt::get(Context, APInt(64, 1));
+    Value *AllocSize = Builder.CreateAdd(Len, One);
+    Value *NewStr = Builder.CreateCall(MallocFunc, AllocSize, "ucase_str");
+    
+    // Loop
+    Function *TheFunction = Builder.GetInsertBlock()->getParent();
+    BasicBlock *LoopBB = BasicBlock::Create(Context, "loop", TheFunction);
+    BasicBlock *AfterBB = BasicBlock::Create(Context, "afterloop", TheFunction);
+    
+    // Init i = 0
+    IRBuilder<> TmpB(&TheFunction->getEntryBlock(), TheFunction->getEntryBlock().begin());
+    AllocaInst *IdxVar = TmpB.CreateAlloca(Type::getInt64Ty(Context), nullptr, "idx");
+    Builder.CreateStore(ConstantInt::get(Context, APInt(64, 0)), IdxVar);
+    
+    Builder.CreateBr(LoopBB);
+    Builder.SetInsertPoint(LoopBB);
+    
+    Value *CurIdx = Builder.CreateLoad(Type::getInt64Ty(Context), IdxVar);
+    Value *Cond = Builder.CreateICmpSLT(CurIdx, Len);
+    
+    BasicBlock *BodyBB = BasicBlock::Create(Context, "body", TheFunction);
+    Builder.CreateCondBr(Cond, BodyBB, AfterBB);
+    
+    Builder.SetInsertPoint(BodyBB);
+
+    Value *SrcPtr = Builder.CreateInBoundsGEP(Type::getInt8Ty(Context), Str, CurIdx);
+    Value *CharVal = Builder.CreateLoad(Type::getInt8Ty(Context), SrcPtr);
+
+    Value *ExtChar = Builder.CreateSExt(CharVal, Type::getInt32Ty(Context));
+    Value *UpperChar = Builder.CreateCall(ToUpperFunc, ExtChar);
+    Value *TruncChar = Builder.CreateTrunc(UpperChar, Type::getInt8Ty(Context));
+
+    Value *DestPtr = Builder.CreateInBoundsGEP(Type::getInt8Ty(Context), NewStr, CurIdx);
+    Builder.CreateStore(TruncChar, DestPtr);
+
+    Value *NextIdx = Builder.CreateAdd(CurIdx, One);
+    Builder.CreateStore(NextIdx, IdxVar);
+    Builder.CreateBr(LoopBB);
+    
+    Builder.SetInsertPoint(AfterBB);
+
+    Value *NullPtr = Builder.CreateInBoundsGEP(Type::getInt8Ty(Context), NewStr, Len);
+    Builder.CreateStore(ConstantInt::get(Type::getInt8Ty(Context), 0), NullPtr);
+    
+    return NewStr;
+}
diff --git a/lib/Lexer/Lexer.cc b/lib/Lexer/Lexer.cc
index d9eb09b..040c3e8 100644
--- a/lib/Lexer/Lexer.cc
+++ b/lib/Lexer/Lexer.cc
@@ -15,6 +15,17 @@ int Lexer::gettok() {
         LastChar = getchar();
     }
 
+    if (LastChar == '"') {
+        StringVal = "";
+        LastChar = getchar();
+        while (LastChar != '"' && LastChar != EOF) {
+            StringVal += LastChar;
+            LastChar = getchar();
+        }
+        if (LastChar == '"') LastChar = getchar();
+        return tok_string_literal;
+    }
+
     if (isalpha(LastChar)) {
         IdentifierStr = LastChar;
         while (isalnum((LastChar = getchar())) || LastChar == '_')
@@ -24,7 +35,8 @@ int Lexer::gettok() {
         if (IdentifierStr == "INTEGER") return tok_integer_kw;
         if (IdentifierStr == "BOOLEAN") return tok_boolean_kw; 
         if (IdentifierStr == "REAL")    return tok_real_kw; 
-        
+        if (IdentifierStr == "STRING")  return tok_string_kw;
+
         if (IdentifierStr == "TRUE") return tok_true;
         if (IdentifierStr == "FALSE") return tok_false;
 
@@ -58,6 +70,12 @@ int Lexer::gettok() {
         if (IdentifierStr == "OR") return tok_or;
         if (IdentifierStr == "NOT") return tok_not;
 
+        if (IdentifierStr == "LENGTH") return tok_length;
+        if (IdentifierStr == "MID") return tok_mid;
+        if (IdentifierStr == "RIGHT") return tok_right;
+        if (IdentifierStr == "LCASE") return tok_lcase;
+        if (IdentifierStr == "UCASE") return tok_ucase;
+
         if (IdentifierStr == "FUNCTION") return tok_function;
         if (IdentifierStr == "ENDFUNCTION") return tok_endfunction;
         if (IdentifierStr == "PROCEDURE") return tok_procedure;
diff --git a/lib/Parser/Parser.cc b/lib/Parser/Parser.cc
index c5e91f9..236aa01 100644
--- a/lib/Parser/Parser.cc
+++ b/lib/Parser/Parser.cc
@@ -95,6 +95,34 @@ std::unique_ptr<ExprAST> Parser::ParseIdentifierExpr() {
     return std::make_unique<VariableExprAST>(IdName);
 }
 
+std::unique_ptr<ExprAST> Parser::ParseStringBuiltin(const std::string &FuncName) {
+    getNextToken();
+    if (CurTok != '(') {
+        fprintf(stderr, "Error: Expected '(' after %s\n", FuncName.c_str());
+        return nullptr;
+    }
+    getNextToken();
+
+    std::vector<std::unique_ptr<ExprAST>> Args;
+    if (CurTok != ')') {
+        while (true) {
+            if (auto Arg = ParseExpression())
+                Args.push_back(std::move(Arg));
+            else
+                return nullptr;
+
+            if (CurTok == ')') break;
+            if (CurTok != ',') {
+                fprintf(stderr, "Error: Expected ',' or ')' in %s\n", FuncName.c_str());
+                return nullptr;
+            }
+            getNextToken();
+        }
+    }
+    getNextToken(); 
+    return std::make_unique<CallExprAST>(FuncName, std::move(Args));
+}
+
 std::unique_ptr<ExprAST> Parser::ParseParenExpr() {
     getNextToken();
     auto V = ParseExpression();
@@ -112,6 +140,17 @@ std::unique_ptr<ExprAST> Parser::ParsePrimary() {
     case tok_identifier: return ParseIdentifierExpr();
     case tok_number_int: return ParseNumberExpr();
     case tok_number_real: return ParseNumberExpr();
+    case tok_string_literal: {
+        auto Res = std::make_unique<StringExprAST>(Lex.StringVal);
+        getNextToken();
+        return Res;
+    }
+    case tok_length: return ParseStringBuiltin("LENGTH");
+    case tok_mid:    return ParseStringBuiltin("MID");
+    case tok_right:  return ParseStringBuiltin("RIGHT");
+    case tok_lcase:  return ParseStringBuiltin("LCASE");
+    case tok_ucase:  return ParseStringBuiltin("UCASE");
+
     case tok_true: {
         getNextToken();
         return std::make_unique<BooleanExprAST>(true);
@@ -241,6 +280,7 @@ std::unique_ptr<StmtAST> Parser::ParseDeclare() {
         if (CurTok == tok_integer_kw) TypeStr = "INTEGER";
         else if (CurTok == tok_real_kw) TypeStr = "REAL";
         else if (CurTok == tok_boolean_kw) TypeStr = "BOOLEAN";
+        else if (CurTok == tok_string_kw) TypeStr = "STRING";
         else {
             fprintf(stderr, "Error: Unknown type in declaration\n");
             return nullptr;