[-Wunsafe-buffer-usage] Check assignments to implicitly immutable count-attributed objects

Check assignments to count-attributed objects that are implicitly
immutable.  For example, assigning to a dependent count that is used in
an inout pointer is not allowed, since the update won't be visible on
the call-site:

```
  void foo(int *__counted_by(count) *out_p, int count) {
    *out_p = ...;
    count = ...; // immutable
  }
```

rdar://161608157

(cherry picked from commit 7f73488)

Author: patrykstefanski

clang/include/clang/Analysis/Analyses/UnsafeBufferUsage.h

@@ -149,6 +149,21 @@ class UnsafeBufferUsageHandler {
                                                      ASTContext &Ctx) {
     handleUnsafeOperation(E, IsRelatedToDecl, Ctx);
   }
+
+  enum class AssignToImmutableObjectKind {
+    PointerToPointer,
+    PointerToDependentCount,
+    PointerDependingOnInoutCount,
+    DependentCountUsedInInoutPointer,
+  };
+
+  virtual void handleAssignToImmutableObject(const BinaryOperator *Assign,
+                                             const ValueDecl *VD,
+                                             AssignToImmutableObjectKind Kind,
+                                             bool IsRelatedToDecl,
+                                             ASTContext &Ctx) {
+    handleUnsafeOperation(Assign, IsRelatedToDecl, Ctx);
+  }
   /* TO_UPSTREAM(BoundsSafety) OFF */
 
   /// Invoked when a fix is suggested against a variable. This function groups

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -14114,6 +14114,13 @@ def warn_assign_to_count_attributed_must_be_simple_stmt : Warning<
   "assignment to %select{count-attributed pointer|dependent count}0 '%1' "
   "must be a simple statement '%1 = ...'">,
   InGroup<UnsafeBufferUsage>, DefaultIgnore;
+def warn_cannot_assign_to_immutable_bounds_attributed_object : Warning<
+  "cannot assign to %select{variable|parameter|member}0 %1 because %select{"
+  "it points to a count-attributed pointer|"
+  "it points to a dependent count|"
+  "its type depends on an inout dependent count|"
+  "it's used as dependent count in an inout count-attributed pointer"
+  "}2">, InGroup<UnsafeBufferUsage>, DefaultIgnore;
 #ifndef NDEBUG
 // Not a user-facing diagnostic. Useful for debugging false negatives in
 // -fsafe-buffer-usage-suggestions (i.e. lack of -Wunsafe-buffer-usage fixits).

clang/lib/Analysis/UnsafeBufferUsage.cpp

@@ -5654,6 +5654,89 @@ struct BoundsAttributedGroupFinder
   }
 };
 
+// Checks if the bounds-attributed group does not assign to implicitly
+// immutable objects (check below which patterns are considered immutable).
+// This function returns false iff at least one assignment modifies
+// bounds-attributed object that should be immutable.
+static bool checkImmutableBoundsAttributedObjects(
+    const BoundsAttributedAssignmentGroup &Group,
+    UnsafeBufferUsageHandler &Handler, ASTContext &Ctx) {
+  bool IsGroupSafe = true;
+
+  for (size_t I = 0, N = Group.Assignments.size(); I < N; ++I) {
+    const BinaryOperator *Assign = Group.Assignments[I];
+    const BoundsAttributedObject &Object = Group.AssignedObjects[I];
+
+    const ValueDecl *VD = Object.Decl;
+    QualType Ty = VD->getType();
+    int DerefLevel = Object.DerefLevel;
+
+    using AssignKind = UnsafeBufferUsageHandler::AssignToImmutableObjectKind;
+
+    // void foo(int *__counted_by(*len) *p, int *len) {
+    //   p = ...;
+    // }
+    if (DerefLevel == 0 && Ty->isPointerType() &&
+        Ty->getPointeeType()->isBoundsAttributedType()) {
+      Handler.handleAssignToImmutableObject(Assign, VD,
+                                            AssignKind::PointerToPointer,
+                                            /*IsRelatedToDecl=*/false, Ctx);
+      IsGroupSafe = false;
+    }
+
+    // void foo(int *__counted_by(*len) *p, int *len) {
+    //   len = ...;
+    // }
+    if (DerefLevel == 0 && VD->isDependentCountWithDeref()) {
+      Handler.handleAssignToImmutableObject(Assign, VD,
+                                            AssignKind::PointerToDependentCount,
+                                            /*IsRelatedToDecl=*/false, Ctx);
+      IsGroupSafe = false;
+    }
+
+    // `p` below should be immutable, because updating `p` would not be visible
+    // on the call-site to `foo`, but `*len` would, which invalidates the
+    // relation between the pointer and its count.
+    // void foo(int *__counted_by(*len) p, int *len) {
+    //   p = ...;
+    // }
+    if (DerefLevel == 0 && Ty->isCountAttributedTypeDependingOnInoutCount()) {
+      Handler.handleAssignToImmutableObject(
+          Assign, VD, AssignKind::PointerDependingOnInoutCount,
+          /*IsRelatedToDecl=*/false, Ctx);
+      IsGroupSafe = false;
+    }
+
+    // Same as above, we cannot update `len`, because it won't be visible on
+    // the call-site.
+    // void foo(int *__counted_by(len) *p, int *__counted_by(len) q, int len) {
+    //   len = ...; // bad because of p
+    // }
+    if (VD->isDependentCountWithoutDeref() &&
+        VD->isDependentCountThatIsUsedInInoutPointer()) {
+      assert(DerefLevel == 0);
+      Handler.handleAssignToImmutableObject(
+          Assign, VD, AssignKind::DependentCountUsedInInoutPointer,
+          /*IsRelatedToDecl=*/false, Ctx);
+      IsGroupSafe = false;
+    }
+  }
+
+  return IsGroupSafe;
+}
+
+// Checks if the bounds-attributed group is safe. This function returns false
+// iff the assignment group is unsafe and diagnostics have been emitted.
+static bool
+checkBoundsAttributedGroup(const BoundsAttributedAssignmentGroup &Group,
+                           UnsafeBufferUsageHandler &Handler, ASTContext &Ctx) {
+  if (!checkImmutableBoundsAttributedObjects(Group, Handler, Ctx))
+    return false;
+
+  // TODO: Add more checks.
+  return true;
+}
+
 static void
 diagnoseTooComplexAssignToBoundsAttributed(const Expr *E, const ValueDecl *VD,
                                            UnsafeBufferUsageHandler &Handler,
@@ -5679,7 +5762,7 @@ static void checkBoundsSafetyAssignments(const Stmt *S,
                                          ASTContext &Ctx) {
   BoundsAttributedGroupFinder Finder(
       [&](const BoundsAttributedAssignmentGroup &Group) {
-        // TODO: Check group constraints.
+        checkBoundsAttributedGroup(Group, Handler, Ctx);
       },
       [&](const Expr *E, const ValueDecl *VD) {
         diagnoseTooComplexAssignToBoundsAttributed(E, VD, Handler, Ctx);

clang/lib/Sema/AnalysisBasedWarnings.cpp

@@ -2609,6 +2609,30 @@ class UnsafeBufferUsageReporter : public UnsafeBufferUsageHandler {
     S.Diag(Loc, diag::warn_assign_to_count_attributed_must_be_simple_stmt)
         << VD->isDependentCount() << VD->getName();
   }
+
+  static int getBoundsAttributedObjectKind(const ValueDecl *VD) {
+    switch (VD->getKind()) {
+    case Decl::Var:
+      return 0;
+    case Decl::ParmVar:
+      return 1;
+    case Decl::Field:
+      return 2;
+    default:
+      llvm_unreachable("unexpected bounds-attributed decl kind");
+      return 0;
+    }
+  }
+
+  void handleAssignToImmutableObject(const BinaryOperator *Assign,
+                                     const ValueDecl *VD,
+                                     AssignToImmutableObjectKind Kind,
+                                     bool IsRelatedToDecl,
+                                     ASTContext &Ctx) override {
+    S.Diag(Assign->getOperatorLoc(),
+           diag::warn_cannot_assign_to_immutable_bounds_attributed_object)
+        << getBoundsAttributedObjectKind(VD) << VD << Kind;
+  }
   /* TO_UPSTREAM(BoundsSafety) OFF */
 
   void handleUnsafeVariableGroup(const VarDecl *Variable,

clang/test/SemaCXX/warn-unsafe-buffer-usage-count-attributed-pointer-assignment.cpp

@@ -69,6 +69,148 @@ void good_struct_self_loop(cb *c) {
   }
 }
 
+// Inout pointer and count
+
+void good_inout_span(int *__counted_by(*count) *p, size_t *count, std::span<int> sp) {
+  *p = sp.data();
+  *count = sp.size();
+}
+
+void bad_inout_span(int *__counted_by(*count) *p, size_t *count, std::span<int> sp) {
+  *p = sp.data(); // TODO-expected-warning{{unsafe assignment to count-attributed pointer}}
+  *count = 42;
+}
+
+void good_inout_subspan_const(int *__counted_by(*count) *p, size_t *count, std::span<int> sp) {
+  *p = sp.first(42).data();
+  *count = 42;
+}
+
+void good_inout_subspan_var(int *__counted_by(*count) *p, size_t *count, std::span<int> sp, size_t new_count) {
+  *p = sp.first(new_count).data();
+  *count = new_count;
+}
+
+void good_inout_subspan_complex(int *__counted_by(*count) *p, size_t *count, std::span<int> sp, size_t i, size_t j) {
+  *p = sp.first(i + j * 2).data();
+  *count = i + j * 2;
+}
+
+void good_inout_span_if(int *__counted_by(*count) *p, size_t *count, std::span<int> sp) {
+  if (p && count) {
+    *p = sp.data();
+    *count = sp.size();
+  }
+}
+
+void bad_inout_span_if(int *__counted_by(*count) *p, size_t *count, std::span<int> sp, size_t size) {
+  if (p && count) {
+    *p = sp.data(); // TODO-expected-warning{{unsafe assignment to count-attributed pointer}}
+    *count = size;
+  }
+}
+
+class inout_class {
+  void good_inout_span(int *__counted_by(*count) *p, size_t *count, std::span<int> sp) {
+    *p = sp.data();
+    *count = sp.size();
+  }
+
+  void bad_inout_span(int *__counted_by(*count) *p, size_t *count, std::span<int> sp) {
+    *p = sp.data(); // TODO-expected-warning{{unsafe assignment to count-attributed pointer}}
+    *count = 42;
+  }
+
+  void good_inout_subspan_const(int *__counted_by(*count) *p, size_t *count, std::span<int> sp) {
+    *p = sp.first(42).data();
+    *count = 42;
+  }
+};
+
+// Inout pointer
+
+void bad_inout_ptr_span(int *__counted_by(count) *p, int count, std::span<int> sp) {
+  *p = sp.data(); // TODO-expected-warning{{unsafe assignment to count-attributed pointer}}
+}
+
+void good_inout_ptr_subspan(int *__counted_by(count) *p, size_t count, std::span<int> sp) {
+  *p = sp.first(count).data();
+}
+
+void good_inout_ptr_const_subspan(int *__counted_by(42) *p, std::span<int> sp) {
+  *p = sp.first(42).data();
+}
+
+void good_inout_ptr_multi_subspan(int *__counted_by(a + b) *p, size_t a, size_t b, std::span<int> sp) {
+  *p = sp.first(a + b).data();
+}
+
+class inout_ptr_class {
+  void bad_inout_ptr_span(int *__counted_by(count) *p, int count, std::span<int> sp) {
+    *p = sp.data(); // TODO-expected-warning{{unsafe assignment to count-attributed pointer}}
+  }
+
+  void good_inout_ptr_subspan(int *__counted_by(count) *p, size_t count, std::span<int> sp) {
+    *p = sp.first(count).data();
+  }
+};
+
+// Immutable pointers/dependent values
+
+void immutable_ptr_to_ptr(int *__counted_by(*count) *p, int *count) {
+  p = nullptr; // expected-warning{{cannot assign to parameter 'p' because it points to a count-attributed pointer}}
+  *count = 0;
+}
+
+void immutable_ptr_to_value(int *__counted_by(*count) *p, int *count) {
+  *p = nullptr;
+  count = nullptr; // expected-warning{{cannot assign to parameter 'count' because it points to a dependent count}}
+}
+
+void immutable_ptr_with_inout_value(int *__counted_by(*count) p, int *count) {
+  p = nullptr; // expected-warning{{cannot assign to parameter 'p' because its type depends on an inout dependent count}}
+  *count = 0;
+}
+
+void immutable_ptr_with_inout_value2(int *__counted_by(*count) p, int *__counted_by(*count) *q, int *count) {
+  p = nullptr;  // expected-warning{{cannot assign to parameter 'p' because its type depends on an inout dependent count}}
+  *q = nullptr;
+  *count = 0;
+}
+
+void immutable_value_with_inout_ptr(int *__counted_by(count) *p, int count) {
+  *p = nullptr;
+  count = 0; // expected-warning{{cannot assign to parameter 'count' because it's used as dependent count in an inout count-attributed pointer}}
+}
+
+void immutable_value_with_inout_ptr2(int *__counted_by(count) p, int *__counted_by(count) *q, int count) {
+  p = nullptr;
+  *q = nullptr;
+  count = 0; // expected-warning{{cannot assign to parameter 'count' because it's used as dependent count in an inout count-attributed pointer}}
+}
+
+class immutable_class {
+  void immutable_ptr_to_ptr(int *__counted_by(*count) *p, int *count) {
+    p = nullptr; // expected-warning{{cannot assign to parameter 'p' because it points to a count-attributed pointer}}
+    *count = 0;
+  }
+
+  void immutable_ptr_to_value(int *__counted_by(*count) *p, int *count) {
+    *p = nullptr;
+    count = nullptr; // expected-warning{{cannot assign to parameter 'count' because it points to a dependent count}}
+  }
+
+  void immutable_ptr_with_inout_value(int *__counted_by(*count) p, int *count) {
+    p = nullptr; // expected-warning{{cannot assign to parameter 'p' because its type depends on an inout dependent count}}
+    *count = 0;
+  }
+
+  void immutable_value_with_inout_ptr(int *__counted_by(count) *p, int count) {
+    *p = nullptr;
+    count = 0; // expected-warning{{cannot assign to parameter 'count' because it's used as dependent count in an inout count-attributed pointer}}
+  }
+};
+
 // Assigns to bounds-attributed that we consider too complex to analyze.
 
 void too_complex_assign_to_ptr(int *__counted_by(count) p, int count, int *q) {