Add spec to FunctionBody

Summary:
Keep a reference to the `SpecBody` of a function in the `FunctionBody`, as a convenience.
The eventual goal is to have `FunctionBody` as a one-stop-shop for anything related to a function.

Reviewed By: robertoaloi

Differential Revision: D55021932

fbshipit-source-id: ae0000b5b5b8601546b7e18e8351892b3b9ccfc9

crates/hir/src/body.rs

@@ -92,6 +92,7 @@ pub struct FunctionBody {
     pub function_id: InFile<FunctionDefId>,
     pub clause_ids: Vec<FunctionClauseId>,
     pub clauses: Arena<Arc<FunctionClauseBody>>,
+    pub spec: Option<Arc<SpecBody>>,
 }
 
 #[derive(Debug, PartialEq, Eq, Clone)]
@@ -425,15 +426,20 @@ impl FunctionBody {
             let mut ctx = lower::Ctx::new(db, BodyOrigin::Invalid(function_id.file_id));
             let name = &fun_def.function_clauses[0].name;
             ctx.set_function_info(name);
-            let (body, source_maps) =
+            let (mut body, source_maps) =
                 ctx.lower_function(function_id, fun_def.function_clause_ids.clone(), &fun_asts);
+            if let Some(spec) = &fun_def.spec {
+                let spec = db.spec_body(InFile::new(spec.file.file_id, spec.spec_id));
+                body.spec = Some(spec);
+            }
             (Arc::new(body), source_maps)
         } else {
             (
                 Arc::new(FunctionBody {
                     function_id,
                     clause_ids: vec![],
                     clauses: Arena::default(),
+                    spec: None,
                 }),
                 vec![],
             )
@@ -446,6 +452,13 @@ impl FunctionBody {
         Some(FormIdx::FunctionClause(*function_id))
     }
 
+    pub fn spec_body(&self) -> Option<&SpecBody> {
+        match &self.spec {
+            Some(spec) => Some(&spec),
+            None => None,
+        }
+    }
+
     pub fn print(&self, db: &dyn InternDatabase, form: &FunctionClause) -> String {
         pretty::print_function_clause(db, self, form)
     }
@@ -620,6 +633,18 @@ impl SpecBody {
         (Arc::new(body), Arc::new(source_map))
     }
 
+    pub fn spec_id(&self) -> Option<SpecId> {
+        if let BodyOrigin::FormIdx {
+            file_id: _,
+            form_id: FormIdx::Spec(spec_id),
+        } = self.body.origin
+        {
+            Some(spec_id)
+        } else {
+            None
+        }
+    }
+
     pub fn print(&self, db: &dyn InternDatabase, form: SpecOrCallback) -> String {
         pretty::print_spec(db, self, form)
     }

crates/hir/src/body/lower.rs

@@ -226,6 +226,7 @@ impl<'a> Ctx<'a> {
                 function_id,
                 clause_ids,
                 clauses,
+                spec: None, // Will be filled in later
             },
             source_maps,
         )