From b35e46265b56e3d11a98e4f4717f79bc6e3026cd Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveth@hclcat.games>
Date: Mon, 5 Dec 2022 19:28:29 +0800
Subject: [PATCH] fix a bug

---
 src/compiler.h |  2 ++
 src/main.cpp   |  4 ++--
 src/pocketpy.h | 44 ++++++++++++++++++++++++++++++++++++++++++++
 tests/math.py  | 10 ++++++++++
 4 files changed, 58 insertions(+), 2 deletions(-)
 create mode 100644 tests/math.py

diff --git a/src/compiler.h b/src/compiler.h
index 2aa71056..7565ee26 100644
--- a/src/compiler.h
+++ b/src/compiler.h
@@ -683,6 +683,7 @@ __LISTCOMP:
         Token tkmodule = compileImportPath();
         consume(TK("import"));
         do {
+            emitCode(OP_DUP_TOP);
             consume(TK("@id"));
             Token tkname = parser->previous;
             int index = getCode()->addName(tkname.str(), NAME_GLOBAL);
@@ -694,6 +695,7 @@ __LISTCOMP:
             index = getCode()->addName(tkname.str(), NAME_GLOBAL);
             emitCode(OP_STORE_NAME_PTR, index);
         } while (match(TK(",")));
+        emitCode(OP_POP_TOP);
         consumeEndStatement();
     }
 
diff --git a/src/main.cpp b/src/main.cpp
index 9301eeef..5ec1146b 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -3,8 +3,8 @@
 
 #include "pocketpy.h"
 
-#define PK_DEBUG_TIME
-//#define PK_DEBUG_THREADED
+//#define PK_DEBUG_TIME
+#define PK_DEBUG_THREADED
 
 struct Timer{
     const char* title;
diff --git a/src/pocketpy.h b/src/pocketpy.h
index 6a59541b..cf5d6935 100644
--- a/src/pocketpy.h
+++ b/src/pocketpy.h
@@ -648,6 +648,49 @@ void __addModuleJson(VM* vm){
     });
 }
 
+void __addModuleMath(VM* vm){
+    PyVar mod = vm->newModule("math");
+    vm->setAttr(mod, "pi", vm->PyFloat(M_PI));
+    vm->setAttr(mod, "e", vm->PyFloat(M_E));
+
+    vm->bindFunc(mod, "log", [](VM* vm, const pkpy::ArgList& args) {
+        vm->__checkArgSize(args, 1);
+        return vm->PyFloat(log(vm->numToFloat(args[0])));
+    });
+
+    vm->bindFunc(mod, "log10", [](VM* vm, const pkpy::ArgList& args) {
+        vm->__checkArgSize(args, 1);
+        return vm->PyFloat(log10(vm->numToFloat(args[0])));
+    });
+
+    vm->bindFunc(mod, "log2", [](VM* vm, const pkpy::ArgList& args) {
+        vm->__checkArgSize(args, 1);
+        return vm->PyFloat(log2(vm->numToFloat(args[0])));
+    });
+
+    vm->bindFunc(mod, "sin", [](VM* vm, const pkpy::ArgList& args) {
+        vm->__checkArgSize(args, 1);
+        return vm->PyFloat(sin(vm->numToFloat(args[0])));
+    });
+
+    vm->bindFunc(mod, "cos", [](VM* vm, const pkpy::ArgList& args) {
+        vm->__checkArgSize(args, 1);
+        return vm->PyFloat(cos(vm->numToFloat(args[0])));
+    });
+
+    vm->bindFunc(mod, "tan", [](VM* vm, const pkpy::ArgList& args) {
+        vm->__checkArgSize(args, 1);
+        return vm->PyFloat(tan(vm->numToFloat(args[0])));
+    });
+
+    vm->bindFunc(mod, "isclose", [](VM* vm, const pkpy::ArgList& args) {
+        vm->__checkArgSize(args, 2);
+        _Float a = vm->numToFloat(args[0]);
+        _Float b = vm->numToFloat(args[1]);
+        return vm->PyBool(fabs(a - b) < 1e-6);
+    });
+}
+
 class _PkExported{
 public:
     virtual ~_PkExported() = default;
@@ -764,6 +807,7 @@ extern "C" {
         __addModuleSys(vm);
         __addModuleTime(vm);
         __addModuleJson(vm);
+        __addModuleMath(vm);
 
         _Code code = compile(vm, __BUILTINS_CODE, "<builtins>");
         if(code == nullptr) exit(1);
diff --git a/tests/math.py b/tests/math.py
new file mode 100644
index 00000000..d43f9761
--- /dev/null
+++ b/tests/math.py
@@ -0,0 +1,10 @@
+from math import log, log10, log2, sin, cos, tan, e, pi, isclose
+
+assert isclose(e, 2.718281828459045)
+assert isclose(pi, 3.141592653589793)
+assert isclose(log(10), 2.302585092994046)
+assert isclose(log10(10), 1.0)
+assert isclose(log2(10), 3.321928094887362)
+assert isclose(sin(0), 0.0)
+assert isclose(cos(0), 1.0)
+assert isclose(tan(0), 0.0)