3 files changed, 48 insertions, 11 deletions

diff --git a/compile.c b/compile.c
index 2f2bc4e3..178a7794 100644
--- a/compile.c
+++ b/compile.c
@@ -520,6 +520,8 @@ CORD compile_statement(env_t *env, ast_t *ast)
     }
     case For: {
         auto for_ = Match(ast, For);
+        // TODO: optimize case for iterating over comprehensions so we don't need to create
+        // an intermediary array/table
         type_t *iter_t = get_type(env, for_->iter);
         env_t *body_scope = for_scope(env, ast);
         loop_ctx_t loop_ctx = (loop_ctx_t){
@@ -1290,7 +1292,13 @@ CORD compile(env_t *env, ast_t *ast)
 
     }
     case Comprehension: {
-        code_err(ast, "Comprehensions cannot be compiled as expressions");
+        ast_t *base = Match(ast, Comprehension)->expr;
+        while (base->tag == Comprehension)
+            base = Match(ast, Comprehension)->expr;
+        if (base->tag == TableEntry)
+            return compile(env, WrapAST(ast, Table, .entries=new(ast_list_t, .ast=ast)));
+        else
+            return compile(env, WrapAST(ast, Array, .items=new(ast_list_t, .ast=ast)));
     }
     case Lambda: {
         auto lambda = Match(ast, Lambda);
diff --git a/parse.c b/parse.c
index 50f3ac2d..2e9787c1 100644
--- a/parse.c
+++ b/parse.c
@@ -50,11 +50,12 @@ static const char *keywords[] = {
 };
 
 enum {NORMAL_FUNCTION=0, EXTERN_FUNCTION=1};
+typedef enum {WHITESPACE_NONE=0, WHITESPACE_SPACES=1, WHITESPACE_NEWLINES=2, WHITESPACE_COMMENTS=4} whitespace_types_e;
 
 static inline size_t some_of(const char **pos, const char *allow);
 static inline size_t some_not(const char **pos, const char *forbid);
 static inline size_t spaces(const char **pos);
-static inline size_t whitespace(const char **pos);
+static inline whitespace_types_e whitespace(const char **pos);
 static inline size_t match(const char **pos, const char *target);
 static inline void expect_str(parse_ctx_t *ctx, const char *start, const char **pos, const char *target, const char *fmt, ...);
 static inline void expect_closing(parse_ctx_t *ctx, const char **pos, const char *target, const char *fmt, ...);
@@ -186,11 +187,18 @@ size_t spaces(const char **pos) {
     return some_of(pos, " \t");
 }
 
-size_t whitespace(const char **pos) {
-    const char *p0 = *pos;
-    while (some_of(pos, " \t\r\n") || comment(pos))
-        continue;
-    return (size_t)(*pos - p0);
+whitespace_types_e whitespace(const char **pos) {
+    whitespace_types_e spaces = WHITESPACE_NONE;
+    for (;;) {
+        if (some_of(pos, " \t")) {
+            spaces |= WHITESPACE_SPACES;
+        } else if (some_of(pos, "\r\n")) {
+            spaces |= WHITESPACE_NEWLINES;
+        } else if (comment(pos)) {
+            spaces |= WHITESPACE_COMMENTS;
+        } else break;
+    }
+    return spaces;
 }
 
 size_t match(const char **pos, const char *target) {
@@ -764,6 +772,12 @@ PARSER(parse_reduction) {
 
     ast_t *iter = optional(ctx, &pos, parse_extended_expr);
     if (!iter) return NULL;
+    ast_t *suffixed = parse_comprehension_suffix(ctx, iter);
+    while (suffixed) {
+        iter = suffixed;
+        pos = suffixed->end;
+        suffixed = parse_comprehension_suffix(ctx, iter);
+    }
 
     ast_t *fallback = NULL;
     if (match_word(&pos, "else"))
@@ -1503,7 +1517,7 @@ PARSER(parse_block) {
             break;
         }
         statements = new(ast_list_t, .ast=stmt, .next=statements);
-        whitespace(&pos);
+        whitespace(&pos); // TODO: check for newline
         if (get_indent(ctx->file, pos) != block_indent) {
             pos = stmt->end; // backtrack
             break;
@@ -1537,7 +1551,11 @@ PARSER(parse_namespace) {
         {
             statements = new(ast_list_t, .ast=stmt, .next=statements);
             pos = stmt->end;
-            whitespace(&pos);
+            whitespace(&pos); // TODO: check for newline
+            // if (!(space_types & WHITESPACE_NEWLINES)) {
+            //     pos = stmt->end;
+            //     break;
+            // }
         } else {
             if (get_indent(ctx->file, next) > indent && next < strchrnul(next, '\n'))
                 parser_err(ctx, next, strchrnul(next, '\n'), "I couldn't parse this namespace statement");
@@ -1828,11 +1846,13 @@ PARSER(parse_doctest) {
     if (match(&pos, "=")) {
         spaces(&pos);
         const char *output_start = pos,
-                   *output_end = strchrnul(pos, '\n');
+                   *output_end = pos + strcspn(pos, "\r\n");
         if (output_end <= output_start)
             parser_err(ctx, output_start, output_end, "You're missing expected output here");
         output = heap_strn(output_start, (size_t)(output_end - output_start));
         pos = output_end;
+    } else {
+        pos = expr->end;
     }
     return NewAST(ctx->file, start, pos, DocTest, .expr=expr, .output=output);
 }
diff --git a/typecheck.c b/typecheck.c
index 2d212f76..ced1f028 100644
--- a/typecheck.c
+++ b/typecheck.c
@@ -374,7 +374,16 @@ type_t *get_type(env_t *env, ast_t *ast)
         code_err(ast, "Table entries should not be typechecked directly");
     }
     case Comprehension: {
-        code_err(ast, "Comprehensions should not be typechecked directly");
+        auto comp = Match(ast, Comprehension);
+        env_t *scope = for_scope(env, FakeAST(For, .iter=comp->iter, .index=comp->key, .value=comp->value));
+        if (comp->expr->tag == Comprehension) {
+            return get_type(scope, comp->expr);
+        } else if (comp->expr->tag == TableEntry) {
+            auto e = Match(comp->expr, TableEntry);
+            return Type(TableType, .key_type=get_type(scope, e->key), .value_type=get_type(scope, e->value));
+        } else {
+            return Type(ArrayType, .item_type=get_type(scope, comp->expr));
+        }
     }
     case FieldAccess: {
         auto access = Match(ast, FieldAccess);