Fix up doctests

5 files changed, 24 insertions, 7 deletions

diff --git a/builtins/functions.c b/builtins/functions.c
index 03dfe738..309ad2ab 100644
--- a/builtins/functions.c
+++ b/builtins/functions.c
@@ -121,7 +121,7 @@ static inline char *without_colors(const char *str)
     return buf;
 }
 
-public void __doctest(CORD label, void *expr, TypeInfo *type, CORD expected, const char *filename, int start, int end)
+public void __doctest(void *expr, TypeInfo *type, CORD expected, const char *filename, int start, int end)
 {
     static file_t *file = NULL;
     if (filename && (file == NULL || strcmp(file->filename, filename) != 0))
@@ -134,7 +134,7 @@ public void __doctest(CORD label, void *expr, TypeInfo *type, CORD expected, con
         CORD expr_str = generic_as_str(expr, USE_COLOR, type);
         CORD type_name = generic_as_str(NULL, false, type);
 
-        CORD_fprintf(stderr, USE_COLOR ? "\x1b[2m%r\x1b[0m %r \x1b[2m: %r\x1b[m\n" : "%r %r : %r\n", label, expr_str, type_name);
+        CORD_fprintf(stderr, USE_COLOR ? "\x1b[2m=\x1b[0m %r \x1b[2m: %r\x1b[m\n" : "= %r : %r\n", expr_str, type_name);
         if (expected) {
             CORD expr_plain = USE_COLOR ? generic_as_str(expr, false, type) : expr_str;
             bool success = (CORD_cmp(expr_str, expected) == 0);
diff --git a/builtins/functions.h b/builtins/functions.h
index 133c76fc..bb8b563b 100644
--- a/builtins/functions.h
+++ b/builtins/functions.h
@@ -10,7 +10,7 @@ extern const char *SSS_HASH_VECTOR;
 
 void fail(CORD fmt, ...);
 CORD builtin_last_err();
-void __doctest(CORD label, void *expr, TypeInfo *type, CORD expected, const char *filename, int start, int end);
+void __doctest(void *expr, TypeInfo *type, CORD expected, const char *filename, int start, int end);
 
 uint32_t generic_hash(const void *obj, const TypeInfo *type);
 int32_t generic_compare(const void *x, const void *y, const TypeInfo *type);
diff --git a/compile.c b/compile.c
index c4c126ad..baebedee 100644
--- a/compile.c
+++ b/compile.c
@@ -360,11 +360,12 @@ CORD compile(env_t *env, ast_t *ast)
     case DocTest: {
         auto test = Match(ast, DocTest);
         CORD src = heap_strn(test->expr->start, (size_t)(test->expr->end - test->expr->start));
+        type_t *expr_t = get_type(env, test->expr);
         if (test->expr->tag == Declare) {
             auto decl = Match(test->expr, Declare);
             return CORD_asprintf(
                 "$var(%r, %r);\n"
-                "__doctest(\"=\", &%r, %r, %r, %r, %ld, %ld);",
+                "__doctest(&%r, %r, %r, %r, %ld, %ld);",
                 compile(env, decl->var), compile(env, decl->value),
                 compile(env, decl->var),
                 compile_type_info(env, get_type(env, decl->value)),
@@ -396,12 +397,22 @@ CORD compile(env_t *env, ast_t *ast)
                 expr_cord,
                 compile(env, WrapAST(test->expr, StringLiteral, .cord=test->output)));
             return CORD_cat(code, "\n}");
+        } else if (expr_t->tag == VoidType || expr_t->tag == AbortType) {
+            return CORD_asprintf(
+                "__doctest((%r, NULL), NULL, NULL, %r, %ld, %ld);",
+                compile(env, test->expr),
+                compile(env, WrapAST(test->expr, StringLiteral, .cord=test->expr->file->filename)),
+                (int64_t)(test->expr->start - test->expr->file->text),
+                (int64_t)(test->expr->end - test->expr->file->text));
         } else {
             return CORD_asprintf(
-                "$test(%r, %r, %r);",
-                compile(env, WrapAST(test->expr, StringLiteral, .cord=src)),
+                "__doctest($stack(%r), %r, %r, %r, %ld, %ld);",
                 compile(env, test->expr),
-                compile(env, WrapAST(test->expr, StringLiteral, .cord=test->output)));
+                compile_type_info(env, expr_t),
+                compile(env, WrapAST(test->expr, StringLiteral, .cord=test->output)),
+                compile(env, WrapAST(test->expr, StringLiteral, .cord=test->expr->file->filename)),
+                (int64_t)(test->expr->start - test->expr->file->text),
+                (int64_t)(test->expr->end - test->expr->file->text));
         }
     }
     case FieldAccess: {
diff --git a/nextlang.c b/nextlang.c
index 99b3c7f8..65500ebb 100644
--- a/nextlang.c
+++ b/nextlang.c
@@ -7,6 +7,7 @@
 #include "ast.h"
 #include "parse.h"
 #include "compile.h"
+#include "typecheck.h"
 #include "types.h"
 
 int main(int argc, char *argv[])
@@ -51,6 +52,7 @@ int main(int argc, char *argv[])
         CORD code = compile_statement(env, stmt->ast);
         if (code)
             CORD_appendf(&env->code->main, "%r\n", code);
+        bind_statement(env, stmt->ast);
     }
 
     CORD fileinfo = CORD_asprintf("#line 0 \"%s\"\n", f->filename);
diff --git a/typecheck.c b/typecheck.c
index a44c60b1..33e2b122 100644
--- a/typecheck.c
+++ b/typecheck.c
@@ -92,6 +92,10 @@ type_t *get_math_type(env_t *env, ast_t *ast, type_t *lhs_t, type_t *rhs_t)
 void bind_statement(env_t *env, ast_t *statement)
 {
     switch (statement->tag) {
+    case DocTest: {
+        bind_statement(env, Match(statement, DocTest)->expr);
+        break;
+    }
     case Declare: {
         auto decl = Match(statement, Declare);
         type_t *type = get_type(env, decl->value);