From dea491575a9dd657c2a5aacbd2611b05899c64db Mon Sep 17 00:00:00 2001
From: Aaron Eline <aaron@correctcomputation.com>
Date: Fri, 18 Dec 2020 12:26:26 -0500
Subject: [PATCH 1/2] CheckedRegions understnad ret types

---
 clang/lib/3C/CheckedRegions.cpp | 22 ++++++++++++++++++++--
 clang/test/3C/checkedregions.c  |  7 +++++++
 clang/test/3C/fn_sets.c         | 10 +++++-----
 3 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/clang/lib/3C/CheckedRegions.cpp b/clang/lib/3C/CheckedRegions.cpp
index 6c9ce210b7d3..34477b8bccb2 100644
--- a/clang/lib/3C/CheckedRegions.cpp
+++ b/clang/lib/3C/CheckedRegions.cpp
@@ -92,7 +92,8 @@ CheckedRegionAdder::findParentCompound(const ast_type_traits::DynTypedNode &N,
   return *Min;
 }
 
-bool CheckedRegionAdder::isFunctionBody(CompoundStmt *S) {
+
+bool isTopLevel(ASTContext *Context, CompoundStmt *S) {
   const auto &Parents = Context->getParents(*S);
   if (Parents.empty()) {
     return false;
@@ -100,6 +101,10 @@ bool CheckedRegionAdder::isFunctionBody(CompoundStmt *S) {
   return Parents[0].get<FunctionDecl>();
 }
 
+bool CheckedRegionAdder::isFunctionBody(CompoundStmt *S) {
+  return isTopLevel(Context, S);
+}
+
 bool CheckedRegionAdder::isParentChecked(
     const ast_type_traits::DynTypedNode &DTN) {
   if (const auto *Parent = findParentCompound(DTN).first) {
@@ -154,13 +159,26 @@ bool CheckedRegionFinder::VisitDoStmt(DoStmt *S) {
 
 bool CheckedRegionFinder::VisitCompoundStmt(CompoundStmt *S) {
   // Visit all subblocks, find all unchecked types
-  bool Localwild = 0;
+  bool Localwild = false;
   for (const auto &SubStmt : S->children()) {
     CheckedRegionFinder Sub(Context, Writer, Info, Seen, Map, EmitWarnings);
     Sub.TraverseStmt(SubStmt);
     Localwild |= Sub.Wild;
   }
 
+  // If we are a function def, need to check return type
+  if (isTopLevel(Context, S)) {
+    const auto &Parents = Context->getParents(*S);
+    assert(!Parents.empty());
+    FunctionDecl* Parent = const_cast<FunctionDecl*>(Parents[0].get<FunctionDecl>());
+    assert(Parent != NULL);
+    auto retType = Parent->getReturnType().getTypePtr();
+    if (retType->isPointerType()) {
+      CVarOption CV = Info.getVariable(Parent, Context);
+      Localwild |= isWild(CV) || containsUncheckedPtr(Parent->getReturnType());
+    }
+  }
+
   markChecked(S, Localwild);
 
   Wild = false;
diff --git a/clang/test/3C/checkedregions.c b/clang/test/3C/checkedregions.c
index 8b3f1700bd8a..992ac010445f 100644
--- a/clang/test/3C/checkedregions.c
+++ b/clang/test/3C/checkedregions.c
@@ -80,3 +80,10 @@ void baz(void) {
     bad();
   }
 }
+
+int* g() { 
+	//CHECK: int* g(void) {
+	return 1;
+}
+
+
diff --git a/clang/test/3C/fn_sets.c b/clang/test/3C/fn_sets.c
index e3c03273d137..036919d4eeac 100644
--- a/clang/test/3C/fn_sets.c
+++ b/clang/test/3C/fn_sets.c
@@ -51,19 +51,19 @@ void foo1(int *z) {
 /* Testing Something with a larger set of functions */
 
 int *a() {
-	//CHECK: int *a(void) _Checked {
+	//CHECK: int *a(void)  {
   return 0;
 }
 int *b() {
-	//CHECK: int *b(void) _Checked {
+	//CHECK: int *b(void)  {
   return 0;
 }
 int *c() {
-	//CHECK: int *c(void) _Checked {
+	//CHECK: int *c(void)  {
   return 0;
 }
 int *d() {
-	//CHECK: int *d(void) _Checked {
+	//CHECK: int *d(void)  {
   return 0;
 }
 int *e() {
@@ -72,7 +72,7 @@ int *e() {
 	//CHECK: return (int*) 1;
 }
 int *i() {
-	//CHECK: _Ptr<int> i(void) _Checked {
+	//CHECK: _Ptr<int> i(void) _Checked  {
   return 0;
 }
 

From 8448c263e1d9612faf74344fcc0120124d00abc1 Mon Sep 17 00:00:00 2001
From: Aaron Eline <aaron@correctcomputation.com>
Date: Fri, 18 Dec 2020 13:14:19 -0500
Subject: [PATCH 2/2] Fixing comments

---
 clang/lib/3C/CheckedRegions.cpp | 21 +++++++++++++--------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/clang/lib/3C/CheckedRegions.cpp b/clang/lib/3C/CheckedRegions.cpp
index 34477b8bccb2..aeb4c40d3532 100644
--- a/clang/lib/3C/CheckedRegions.cpp
+++ b/clang/lib/3C/CheckedRegions.cpp
@@ -26,6 +26,18 @@
 using namespace llvm;
 using namespace clang;
 
+
+// Check if the compound statement is a function body
+// Used in both visitors so abstracted to a function
+bool isTopLevel(ASTContext *Context, CompoundStmt *S) {
+  const auto &Parents = Context->getParents(*S);
+  if (Parents.empty()) {
+    return false;
+  }
+  // Ensure that our parent is a functiondecl
+  return Parents[0].get<FunctionDecl>() != nullptr;
+}
+
 // CheckedRegionAdder
 
 bool CheckedRegionAdder::VisitCompoundStmt(CompoundStmt *S) {
@@ -93,13 +105,6 @@ CheckedRegionAdder::findParentCompound(const ast_type_traits::DynTypedNode &N,
 }
 
 
-bool isTopLevel(ASTContext *Context, CompoundStmt *S) {
-  const auto &Parents = Context->getParents(*S);
-  if (Parents.empty()) {
-    return false;
-  }
-  return Parents[0].get<FunctionDecl>();
-}
 
 bool CheckedRegionAdder::isFunctionBody(CompoundStmt *S) {
   return isTopLevel(Context, S);
@@ -171,7 +176,7 @@ bool CheckedRegionFinder::VisitCompoundStmt(CompoundStmt *S) {
     const auto &Parents = Context->getParents(*S);
     assert(!Parents.empty());
     FunctionDecl* Parent = const_cast<FunctionDecl*>(Parents[0].get<FunctionDecl>());
-    assert(Parent != NULL);
+    assert(Parent != nullptr);
     auto retType = Parent->getReturnType().getTypePtr();
     if (retType->isPointerType()) {
       CVarOption CV = Info.getVariable(Parent, Context);