5 files changed, 163 insertions, 2 deletions

diff --git a/crates/ra_hir/src/code_model.rs b/crates/ra_hir/src/code_model.rs
index a379b9f49..131180a63 100644
--- a/crates/ra_hir/src/code_model.rs
+++ b/crates/ra_hir/src/code_model.rs
@@ -36,6 +36,7 @@ use rustc_hash::FxHashSet;
 
 use crate::{
     db::{DefDatabase, HirDatabase},
+    diagnostics::UnsafeValidator,
     has_source::HasSource,
     CallableDef, HirDisplay, InFile, Name,
 };
@@ -677,7 +678,9 @@ impl Function {
         let _p = profile("Function::diagnostics");
         let infer = db.infer(self.id.into());
         infer.add_diagnostics(db, self.id, sink);
-        let mut validator = ExprValidator::new(self.id, infer, sink);
+        let mut validator = ExprValidator::new(self.id, infer.clone(), sink);
+        validator.validate_body(db);
+        let mut validator = UnsafeValidator::new(&self, infer, sink);
         validator.validate_body(db);
     }
 }
diff --git a/crates/ra_hir/src/diagnostics.rs b/crates/ra_hir/src/diagnostics.rs
index c82883d0c..562f3fe5c 100644
--- a/crates/ra_hir/src/diagnostics.rs
+++ b/crates/ra_hir/src/diagnostics.rs
@@ -2,3 +2,53 @@
 pub use hir_def::diagnostics::UnresolvedModule;
 pub use hir_expand::diagnostics::{AstDiagnostic, Diagnostic, DiagnosticSink};
 pub use hir_ty::diagnostics::{MissingFields, MissingMatchArms, MissingOkInTailExpr, NoSuchField};
+
+use std::sync::Arc;
+
+use crate::code_model::Function;
+use crate::db::HirDatabase;
+use crate::has_source::HasSource;
+use hir_ty::{
+    diagnostics::{MissingUnsafe, UnnecessaryUnsafe},
+    expr::unsafe_expressions,
+    InferenceResult,
+};
+use ra_syntax::AstPtr;
+
+pub struct UnsafeValidator<'a, 'b: 'a> {
+    func: &'a Function,
+    infer: Arc<InferenceResult>,
+    sink: &'a mut DiagnosticSink<'b>,
+}
+
+impl<'a, 'b> UnsafeValidator<'a, 'b> {
+    pub fn new(
+        func: &'a Function,
+        infer: Arc<InferenceResult>,
+        sink: &'a mut DiagnosticSink<'b>,
+    ) -> UnsafeValidator<'a, 'b> {
+        UnsafeValidator { func, infer, sink }
+    }
+
+    pub fn validate_body(&mut self, db: &dyn HirDatabase) {
+        let def = self.func.id.into();
+        let unsafe_expressions = unsafe_expressions(db, self.infer.as_ref(), def);
+        let func_data = db.function_data(self.func.id);
+        let unnecessary = func_data.is_unsafe && unsafe_expressions.len() == 0;
+        let missing = !func_data.is_unsafe && unsafe_expressions.len() > 0;
+        if !(unnecessary || missing) {
+            return;
+        }
+
+        let in_file = self.func.source(db);
+        let file = in_file.file_id;
+        let fn_def = AstPtr::new(&in_file.value);
+        let fn_name = func_data.name.clone().into();
+
+        if unnecessary {
+            self.sink.push(UnnecessaryUnsafe { file, fn_def, fn_name })
+        } else {
+            self.sink.push(MissingUnsafe { file, fn_def, fn_name })
+        }
+    }
+}
diff --git a/crates/ra_hir_ty/src/diagnostics.rs b/crates/ra_hir_ty/src/diagnostics.rs
index ebd9cb08f..3469cc680 100644
--- a/crates/ra_hir_ty/src/diagnostics.rs
+++ b/crates/ra_hir_ty/src/diagnostics.rs
@@ -169,3 +169,61 @@ impl AstDiagnostic for BreakOutsideOfLoop {
         ast::Expr::cast(node).unwrap()
     }
 }
+
+#[derive(Debug)]
+pub struct MissingUnsafe {
+    pub file: HirFileId,
+    pub fn_def: AstPtr<ast::FnDef>,
+    pub fn_name: Name,
+}
+
+impl Diagnostic for MissingUnsafe {
+    fn message(&self) -> String {
+        format!("Missing unsafe marker on fn `{}`", self.fn_name)
+    }
+    fn source(&self) -> InFile<SyntaxNodePtr> {
+        InFile { file_id: self.file, value: self.fn_def.clone().into() }
+    }
+    fn as_any(&self) -> &(dyn Any + Send + 'static) {
+        self
+    }
+}
+
+impl AstDiagnostic for MissingUnsafe {
+    type AST = ast::FnDef;
+
+    fn ast(&self, db: &impl AstDatabase) -> Self::AST {
+        let root = db.parse_or_expand(self.source().file_id).unwrap();
+        let node = self.source().value.to_node(&root);
+        ast::FnDef::cast(node).unwrap()
+    }
+}
+
+#[derive(Debug)]
+pub struct UnnecessaryUnsafe {
+    pub file: HirFileId,
+    pub fn_def: AstPtr<ast::FnDef>,
+    pub fn_name: Name,
+}
+
+impl Diagnostic for UnnecessaryUnsafe {
+    fn message(&self) -> String {
+        format!("Unnecessary unsafe marker on fn `{}`", self.fn_name)
+    }
+    fn source(&self) -> InFile<SyntaxNodePtr> {
+        InFile { file_id: self.file, value: self.fn_def.clone().into() }
+    }
+    fn as_any(&self) -> &(dyn Any + Send + 'static) {
+        self
+    }
+}
+
+impl AstDiagnostic for UnnecessaryUnsafe {
+    type AST = ast::FnDef;
+
+    fn ast(&self, db: &impl AstDatabase) -> Self::AST {
+        let root = db.parse_or_expand(self.source().file_id).unwrap();
+        let node = self.source().value.to_node(&root);
+        ast::FnDef::cast(node).unwrap()
+    }
+}
diff --git a/crates/ra_hir_ty/src/expr.rs b/crates/ra_hir_ty/src/expr.rs
index 7db928dde..795f1762c 100644
--- a/crates/ra_hir_ty/src/expr.rs
+++ b/crates/ra_hir_ty/src/expr.rs
@@ -2,7 +2,7 @@
 
 use std::sync::Arc;
 
-use hir_def::{path::path, resolver::HasResolver, AdtId, FunctionId};
+use hir_def::{path::path, resolver::HasResolver, AdtId, DefWithBodyId, FunctionId};
 use hir_expand::diagnostics::DiagnosticSink;
 use ra_syntax::{ast, AstPtr};
 use rustc_hash::FxHashSet;
@@ -312,3 +312,25 @@ pub fn record_pattern_missing_fields(
     }
     Some((variant_def, missed_fields, exhaustive))
 }
+
+pub fn unsafe_expressions(
+    db: &dyn HirDatabase,
+    infer: &InferenceResult,
+    def: DefWithBodyId,
+) -> Vec<ExprId> {
+    let mut unsafe_expr_ids = vec![];
+    let body = db.body(def);
+    for (id, expr) in body.exprs.iter() {
+        if let Expr::Call { callee, .. } = expr {
+            if infer
+                .method_resolution(*callee)
+                .map(|func| db.function_data(func).is_unsafe)
+                .unwrap_or(false)
+            {
+                unsafe_expr_ids.push(id);
+            }
+        }
+    }
+
+    unsafe_expr_ids
+}
diff --git a/crates/ra_ide/src/syntax_highlighting/tests.rs b/crates/ra_ide/src/syntax_highlighting/tests.rs
index b7fad9719..39cd74ac3 100644
--- a/crates/ra_ide/src/syntax_highlighting/tests.rs
+++ b/crates/ra_ide/src/syntax_highlighting/tests.rs
@@ -370,3 +370,31 @@ fn check_highlighting(ra_fixture: &str, snapshot: &str, rainbow: bool) {
     fs::write(dst_file, &actual_html).unwrap();
     assert_eq_text!(expected_html, actual_html);
 }
+
+#[test]
+fn test_unsafe_highlighting() {
+    let (analysis, file_id) = single_file(
+        r#"
+unsafe fn unsafe_fn() {}
+
+struct HasUnsafeFn;
+
+impl HasUnsafeFn {
+    unsafe fn unsafe_method(&self) {}
+}
+
+fn main() {
+    unsafe {
+        unsafe_fn();
+        HasUnsafeFn.unsafe_method();
+    }
+}
+"#
+        .trim(),
+    );
+    let dst_file = project_dir().join("crates/ra_ide/src/snapshots/highlight_unsafe.html");
+    let actual_html = &analysis.highlight_as_html(file_id, false).unwrap();
+    let expected_html = &read_text(&dst_file);
+    fs::write(dst_file, &actual_html).unwrap();
+    assert_eq_text!(expected_html, actual_html);
+}