Add Matplotlib plot support

codon/app/main.cpp

@@ -1,6 +1,7 @@
 #include <algorithm>
 #include <cstdio>
 #include <cstdlib>
+#include <fstream>
 #include <iostream>
 #include <sstream>
 #include <string>
@@ -226,17 +227,18 @@ int jitMode(const std::vector<const char *> &args) {
   llvm::cantFail(jit.init());
   fmt::print(">>> Codon JIT v{} <<<\n", CODON_VERSION);
   std::string code;
+  // std::ifstream qq("scratch.codon");
   for (std::string line; std::getline(std::cin, line);) {
     if (line != "#%%") {
       code += line + "\n";
     } else {
-      fmt::print("{}\n\n[done]\n\n", jitExec(&jit, code));
+      fmt::print("{}[done]\n", jitExec(&jit, code));
       code = "";
       fflush(stdout);
     }
   }
   if (!code.empty())
-    fmt::print("{}\n\n[done]\n\n", jitExec(&jit, code));
+    fmt::print("{}[done]\n", jitExec(&jit, code));
   return EXIT_SUCCESS;
 }
 

codon/sir/analyze/dataflow/cfg.h

@@ -481,6 +481,7 @@ public:
   void visit(const dsl::CustomInstr *v) override;
 
   template <typename NodeType> void process(const NodeType *v) {
+    if (!v) return;
     if (seenIds.find(v->getId()) != seenIds.end())
       return;
     seenIds.insert(v->getId());

codon/sir/util/operator.h

@@ -68,8 +68,10 @@ public:
   }
 
   void visit(BodiedFunc *f) override {
-    seen.insert(f->getBody()->getId());
-    process(f->getBody());
+    if (f->getBody()) {
+      seen.insert(f->getBody()->getId());
+      process(f->getBody());
+    }
   }
 
   LAMBDA_VISIT(VarValue);

stdlib/internal/builtin.codon

@@ -331,7 +331,7 @@ class int:
         return result
 
 
-def _jit_display(x, s: Static[str], bundle: Set[str] = set()):
+def _jit_display(x, s: Static[str], bundle: Set[str] = Set[str]()):
     if hasattr(x, "_repr_mimebundle_") and s == "jupyter":
         d = x._repr_mimebundle_(bundle)
         # TODO: pick appropriate mime

stdlib/internal/python.codon

@@ -35,9 +35,32 @@ PyDict_SetItem = Function[[cobj, cobj, cobj], cobj](cobj())
 PyDict_Next = Function[[cobj, Ptr[int], Ptr[cobj], Ptr[cobj]], int](cobj())
 PyObject_GetIter = Function[[cobj], cobj](cobj())
 PyIter_Next = Function[[cobj], cobj](cobj())
+PyObject_HasAttrString = Function[[cobj, cobj], int](cobj())
+PyImport_AddModule = Function[[cobj], cobj](cobj())
 
 _PY_MODULE_CACHE = Dict[str, pyobj]()
 
+_PY_INIT = """
+import io
+
+cls = None
+try:
+     import matplotlib.figure
+     cls = matplotlib.figure.Figure
+except ModuleNotFoundError:
+     pass
+
+def __codon_repr__(fig):
+     if cls and isinstance(fig, cls):
+          stream = io.StringIO()
+          fig.savefig(stream, format="svg")
+          return 'image/svg+xml', stream.getvalue()
+     elif hasattr(fig, "_repr_html_"):
+          return 'text/html', fig._repr_html_()
+     else:
+          return 'text/plain', fig.__repr__()
+"""
+
 _PY_INITIALIZED = False
 def init():
     global _PY_INITIALIZED
@@ -115,8 +138,14 @@ def init():
     PyObject_GetIter = dlsym(hnd, "PyObject_GetIter")
     global PyIter_Next
     PyIter_Next = dlsym(hnd, "PyIter_Next")
+    global PyObject_HasAttrString
+    PyObject_HasAttrString = dlsym(hnd, "PyObject_HasAttrString")
+    global PyImport_AddModule
+    PyImport_AddModule = dlsym(hnd, "PyImport_AddModule")
+
 
     Py_Initialize()
+    PyRun_SimpleString(_PY_INIT.c_str())
     _PY_INITIALIZED = True
 
 def ensure_initialized():
@@ -229,20 +258,19 @@ class pyobj:
     def get(self, T: type) -> T:
         return T.__from_py__(self)
 
+    def _main_module():
+        m = PyImport_AddModule("__main__".c_str())
+        return pyobj(m)
+
     def _repr_mimebundle_(self, bundle = Set[str]()):
-        # p = PyObject_GetAttrString(self.p, '_repr_mimebundle_'.c_str())
-        # if p != cobj():
-        #     return Dict[str, str].__from_py__(self._getattr("_repr_mimebundle_").__call__())
-        # else:
-        p = PyObject_GetAttrString(self.p, '_repr_html_'.c_str())
-        if p != cobj():
-            return {'text/html': str.__from_py__(self._getattr("_repr_html_").__call__())}
-        return {'text/plain': self.__repr__()}
+        fn = pyobj._main_module()._getattr("__codon_repr__")
+        assert fn.p != cobj(), "cannot find python.__codon_repr__"
+        mime, txt = fn.__call__(self).get(Tuple[str, str])
+        return {mime: txt}
 
 def none():
     raise NotImplementedError()
 
-
 # Type conversions
 
 def py(x) -> pyobj: