[-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

Author: patrykstefanski

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

@@ -150,6 +150,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

@@ -14267,6 +14267,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

@@ -5818,6 +5818,83 @@ struct BoundsAttributedGroupFinder
   }
 };
 
+static bool checkImmutableBoundsAttributedObjects(
+    const BoundsAttributedAssignmentGroup &Group,
+    UnsafeBufferUsageHandler &Handler, ASTContext &Ctx) {
+  bool Ok = 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);
+      Ok = 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);
+      Ok = 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);
+      Ok = 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);
+      Ok = false;
+    }
+  }
+
+  return Ok;
+}
+
+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,
@@ -5843,7 +5920,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

@@ -2619,6 +2619,24 @@ 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) {
+    if (isa<ParmVarDecl>(VD))
+      return 1;
+    if (isa<FieldDecl>(VD))
+      return 2;
+    return 0; // variable
+  }
+
+  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,94 @@ 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 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;
+  }
+}
+
+// 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();
+}
+
+// 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}}
+}
+
 // 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) {