3 files changed, 11 insertions, 3 deletions
diff --git a/lib/CodeGen/ArithmeticHandler.cc b/lib/CodeGen/ArithmeticHandler.cc
index fa71e0c..b7f0665 100644
--- a/lib/CodeGen/ArithmeticHandler.cc
+++ b/lib/CodeGen/ArithmeticHandler.cc
@@ -1,5 +1,6 @@
 #include "cps/ArithmeticHandler.h"
 #include "cps/Lexer.h"
+#include "cps/RuntimeCheck.h"
 #include <cstdio>
 
 using namespace llvm;
@@ -32,7 +33,7 @@ Value *ArithmeticHandler::emitBinaryOp(int Op, Value *LHS, Value *RHS, int Line)
             return nullptr;
         }
 
-        // TODO: DIV divide by 0 runtime check
+        RuntimeChecker.emitDivZeroCheck(RVal, Line);
         return Builder.CreateFDiv(LVal, RVal, "divtmp");
     }
 
@@ -42,6 +43,7 @@ Value *ArithmeticHandler::emitBinaryOp(int Op, Value *LHS, Value *RHS, int Line)
             HadError = true;
             return nullptr;
         }
+        RuntimeChecker.emitDivZeroCheck(RHS, Line);
         if (Op == tok_div) return Builder.CreateSDiv(LHS, RHS, "div_int_tmp");
         if (Op == tok_mod) return Builder.CreateSRem(LHS, RHS, "mod_tmp");
     }
diff --git a/lib/CodeGen/CodeGen.cc b/lib/CodeGen/CodeGen.cc
index a6669a6..d5942e9 100644
--- a/lib/CodeGen/CodeGen.cc
+++ b/lib/CodeGen/CodeGen.cc
@@ -85,7 +85,8 @@ CodeGen::CodeGen() {
                                             Symbols,
                                             HadError);
 
-    ArithHandler = std::make_unique<ArithmeticHandler>(*TheContext, *Builder, HadError);
+    ArithHandler = std::make_unique<ArithmeticHandler>(*TheContext, *Builder, HadError,
+                                                     *RuntimeChecker);
     StrHandler = std::make_unique<StringHandler>(*TheContext, *Builder, *TheModule);
     StrConvHandler = std::make_unique<StringConversionHandler>(*TheContext, *Builder, *TheModule);
     Files = std::make_unique<FileHandler>(*TheContext, *Builder, *TheModule, *RuntimeChecker);
diff --git a/lib/CodeGen/RuntimeCheck.cc b/lib/CodeGen/RuntimeCheck.cc
index 40304fd..8075873 100644
--- a/lib/CodeGen/RuntimeCheck.cc
+++ b/lib/CodeGen/RuntimeCheck.cc
@@ -57,7 +57,12 @@ void RuntimeCheck::emitErrorAndExit(Value *Condition, Value *Msg, ArrayRef<Value
 }
 
 void RuntimeCheck::emitDivZeroCheck(Value *Divisor, int Line) {
-    Value *IsZero = Builder.CreateICmpEQ(Divisor, ConstantInt::get(TheContext, APInt(64, 0)), "is_zero");
+    Value *IsZero;
+    if (Divisor->getType()->isDoubleTy()) {
+        IsZero = Builder.CreateFCmpOEQ(Divisor, ConstantFP::get(Divisor->getType(), 0.0), "is_zero");
+    } else {
+        IsZero = Builder.CreateICmpEQ(Divisor, ConstantInt::get(Divisor->getType(), 0), "is_zero");
+    }
     emitErrorAndExit(IsZero, DivZeroMsg, Line);
 }