handle AnnotatedBody everywhere. It compiles now!

compiler/can/src/def.rs

@@ -139,77 +139,37 @@ pub fn canonicalize_defs<'a>(
     let mut refs_by_symbol = MutMap::default();
     let mut can_defs_by_symbol = HashMap::with_capacity_and_hasher(num_defs, default_hasher());
     let mut pending = Vec::with_capacity(num_defs); // TODO bump allocate this!
-    let mut iter = loc_defs.iter().peekable();
 
     // Canonicalize all the patterns, record shadowing problems, and store
     // the ast::Expr values in pending_exprs for further canonicalization
     // once we've finished assembling the entire scope.
-    while let Some(loc_def) = iter.next() {
+    for loc_def in loc_defs {
-        // Any time we have an Annotation followed immediately by a Body,
+        match to_pending_def(env, var_store, &loc_def.value, &mut scope, pattern_type) {
-        // check to see if their patterns are equivalent. If they are,
+            None => (),
-        // turn it into a TypedBody. Otherwise, give an error.
+            Some((new_output, pending_def)) => {
-        let (new_output, pending_def) = match &loc_def.value {
-            Annotation(pattern, annotation) | Nested(Annotation(pattern, annotation)) => {
-                match iter.peek() {
-                    Some(Located {
-                        value: Body(body_pattern, body_expr),
-                        region: body_region,
-                    }) => {
-                        if pattern.value.equivalent(&body_pattern.value) {
-                            iter.next();
-
-                            pending_typed_body(
-                                env,
-                                body_pattern,
-                                annotation,
-                                body_expr,
-                                var_store,
-                                &mut scope,
-                                pattern_type,
-                            )
-                        } else if loc_def.region.lines_between(body_region) > 1 {
-                            // there is a line of whitespace between the annotation and the body
-                            // treat annotation and body separately
-                            to_pending_def(env, var_store, &loc_def.value, &mut scope, pattern_type)
-                        } else {
-                            // the pattern of the annotation does not match the pattern of the body directly below it
-                            env.problems.push(Problem::SignatureDefMismatch {
-                                annotation_pattern: pattern.region,
-                                def_pattern: body_pattern.region,
-                            });
-
-                            // both the annotation and definition are skipped!
-                            iter.next();
-                            continue;
-                        }
-                    }
-                    _ => to_pending_def(env, var_store, &loc_def.value, &mut scope, pattern_type),
-                }
-            }
-            _ => to_pending_def(env, var_store, &loc_def.value, &mut scope, pattern_type),
-        };
-
-        output.union(new_output);
-
                 // store the top-level defs, used to ensure that closures won't capture them
                 if let PatternType::TopLevelDef = pattern_type {
                     match &pending_def {
                         PendingDef::AnnotationOnly(_, loc_can_pattern, _)
                         | PendingDef::Body(_, loc_can_pattern, _)
-                | PendingDef::TypedBody(_, loc_can_pattern, _, _) => env.top_level_symbols.extend(
+                        | PendingDef::TypedBody(_, loc_can_pattern, _, _) => {
+                            env.top_level_symbols.extend(
                                 bindings_from_patterns(std::iter::once(loc_can_pattern))
                                     .iter()
                                     .map(|t| t.0),
-                ),
+                            )
+                        }
                         PendingDef::Alias { .. } | PendingDef::InvalidAlias => {}
                     }
                 }
-
                 // Record the ast::Expr for later. We'll do another pass through these
                 // once we have the entire scope assembled. If we were to canonicalize
                 // the exprs right now, they wouldn't have symbols in scope from defs
                 // that get would have gotten added later in the defs list!
                 pending.push(pending_def);
+                output.union(new_output);
+            }
+        }
     }
 
     if cfg!(debug_assertions) {
@@ -1344,7 +1304,7 @@ fn to_pending_def<'a>(
     def: &'a ast::Def<'a>,
     scope: &mut Scope,
     pattern_type: PatternType,
-) -> (Output, PendingDef<'a>) {
+) -> Option<(Output, PendingDef<'a>)> {
     use roc_parse::ast::Def::*;
 
     match def {
@@ -1359,10 +1319,10 @@ fn to_pending_def<'a>(
                 loc_pattern.region,
             );
 
-            (
+            Some((
                 output,
                 PendingDef::AnnotationOnly(loc_pattern, loc_can_pattern, loc_ann),
-            )
+            ))
         }
         Body(loc_pattern, loc_expr) => {
             // This takes care of checking for shadowing and adding idents to scope.
@@ -1375,10 +1335,41 @@ fn to_pending_def<'a>(
                 loc_pattern.region,
             );
 
-            (
+            Some((
                 output,
                 PendingDef::Body(loc_pattern, loc_can_pattern, loc_expr),
-            )
+            ))
+        }
+
+        AnnotatedBody {
+            ann_pattern,
+            ann_type,
+            comment: _,
+            body_pattern,
+            body_expr,
+        } => {
+            if ann_pattern.value.equivalent(&body_pattern.value) {
+                Some(pending_typed_body(
+                    env,
+                    ann_pattern,
+                    ann_type,
+                    body_expr,
+                    var_store,
+                    scope,
+                    pattern_type,
+                ))
+            } else {
+                // the pattern of the annotation does not match the pattern of the body direc
+                env.problems.push(Problem::SignatureDefMismatch {
+                    annotation_pattern: ann_pattern.region,
+                    def_pattern: body_pattern.region,
+                });
+
+                // TODO: Should we instead build some PendingDef::InvalidAnnotatedBody ? This would
+                // remove the `Option` on this function (and be probably more reliable for further
+                // problem/error reporting)
+                None
+            }
         }
 
         Alias { name, vars, ann } => {
@@ -1411,12 +1402,12 @@ fn to_pending_def<'a>(
                                     region: loc_var.region,
                                 });
 
-                                return (Output::default(), PendingDef::InvalidAlias);
+                                return Some((Output::default(), PendingDef::InvalidAlias));
                             }
                         }
                     }
 
-                    (
+                    Some((
                         Output::default(),
                         PendingDef::Alias {
                             name: Located {
@@ -1426,7 +1417,7 @@ fn to_pending_def<'a>(
                             vars: can_rigids,
                             ann,
                         },
-                    )
+                    ))
                 }
 
                 Err((original_region, loc_shadowed_symbol)) => {
@@ -1435,7 +1426,7 @@ fn to_pending_def<'a>(
                         shadow: loc_shadowed_symbol,
                     });
 
-                    (Output::default(), PendingDef::InvalidAlias)
+                    Some((Output::default(), PendingDef::InvalidAlias))
                 }
             }
         }

compiler/can/src/operator.rs

@@ -47,6 +47,8 @@ pub fn desugar_def<'a>(arena: &'a Bump, def: &'a Def<'a>) -> Def<'a> {
         Nested(alias @ Alias { .. }) => Nested(alias),
         ann @ Annotation(_, _) => Nested(ann),
         Nested(ann @ Annotation(_, _)) => Nested(ann),
+        ann_body @ AnnotatedBody { .. } => Nested(ann_body),
+        Nested(ann_body @ AnnotatedBody { .. }) => Nested(ann_body),
         Nested(NotYetImplemented(s)) => todo!("{}", s),
         NotYetImplemented(s) => todo!("{}", s),
     }

compiler/fmt/src/def.rs

@@ -15,7 +15,7 @@ impl<'a> Formattable<'a> for Def<'a> {
                 loc_pattern.is_multiline() || loc_annotation.is_multiline()
             }
             Body(loc_pattern, loc_expr) => loc_pattern.is_multiline() || loc_expr.is_multiline(),
-
+            AnnotatedBody { .. } => true, // Sebbes: not really sure here...
             SpaceBefore(sub_def, spaces) | SpaceAfter(sub_def, spaces) => {
                 spaces.iter().any(|s| is_comment(s)) || sub_def.is_multiline()
             }
@@ -58,6 +58,24 @@ impl<'a> Formattable<'a> for Def<'a> {
             Body(loc_pattern, loc_expr) => {
                 fmt_body(buf, &loc_pattern.value, &loc_expr.value, indent);
             }
+            AnnotatedBody {
+                ann_pattern,
+                ann_type,
+                comment,
+                body_pattern,
+                body_expr,
+            } => {
+                ann_pattern.format(buf, indent);
+                buf.push_str(" : ");
+                ann_type.format(buf, indent);
+                if let Some(comment_str) = comment {
+                    buf.push_str(" # ");
+                    buf.push_str(comment_str.trim());
+                }
+                buf.push_str("\n");
+                fmt_body(buf, &body_pattern.value, &body_expr.value, indent);
+            }
+
             SpaceBefore(sub_def, spaces) => {
                 fmt_spaces(buf, spaces.iter(), indent);
                 sub_def.format(buf, indent);

compiler/load/src/docs.rs

@@ -92,6 +92,25 @@ fn generate_module_doc<'a>(
             _ => (acc, None),
         },
 
+        AnnotatedBody{ ann_pattern, .. } => match ann_pattern.value {
+            Pattern::Identifier(identifier) => {
+                // Check if the definition is exposed
+                if exposed_ident_ids
+                    .get_id(&InlinableString::from(identifier))
+                    .is_some()
+                {
+                    let entry = DocEntry {
+                        name: identifier.to_string(),
+                        docs: before_comments_or_new_lines.and_then(comments_or_new_lines_to_docs),
+                    };
+                    acc.push(entry);
+                }
+                (acc, None)
+            }
+
+            _ => (acc, None),
+        },
+
         Alias {
             name: _,
             vars: _,

compiler/parse/src/ast.rs

@@ -281,8 +281,8 @@ pub enum Def<'a> {
         ann_pattern: Loc<Pattern<'a>>,
         ann_type: Loc<TypeAnnotation<'a>>,
         comment: Option<&'a str>,
-        body_pattern: Loc<Pattern<'a>>,
+        body_pattern: &'a Loc<Pattern<'a>>,
-        body_expr: Loc<Expr<'a>>,
+        body_expr: &'a Loc<Expr<'a>>,
     },
 
     // Blank Space (e.g. comments, spaces, newlines) before or after a def.

compiler/parse/src/expr.rs

@@ -509,12 +509,9 @@ fn to_def<'a>(
         Either::First((body_pattern, body_expr)) => {
             Def::Body(arena.alloc(body_pattern), arena.alloc(body_expr))
         }
-        Either::Second(((ann_pattern, ann_type), None)) => annotation_or_alias(
+        Either::Second(((ann_pattern, ann_type), None)) => {
-            arena,
+            annotation_or_alias(arena, &ann_pattern.value, ann_pattern.region, ann_type)
-            &ann_pattern.value,
+        }
-            ann_pattern.region,
-            ann_type,
-        ),
         Either::Second((
             (ann_pattern, ann_type),
             Some((opt_comment, (body_pattern, body_expr))),
@@ -522,8 +519,8 @@ fn to_def<'a>(
             ann_pattern: ann_pattern,
             ann_type: ann_type,
             comment: opt_comment,
-            body_pattern: body_pattern,
+            body_pattern: arena.alloc(body_pattern),
-            body_expr: body_expr,
+            body_expr: arena.alloc(body_expr),
         },
     }
 }