[V2][PATCH 2/3] Use the counted_by atribute info in builtin object size [PR108896]

[Qing Zhao <qing.zhao@oracle.com>]

gcc/ChangeLog:

	PR C/108896
	* tree-object-size.cc (addr_object_size): Use the counted_by
	attribute info.
	* tree.cc (component_ref_has_counted_by_p): New function.
	(component_ref_get_counted_by): New function.
	* tree.h (component_ref_has_counted_by_p): New prototype.
	(component_ref_get_counted_by): New prototype.

gcc/testsuite/ChangeLog:

	PR C/108896
	* gcc.dg/flex-array-counted-by-2.c: New test.
	* gcc.dg/flex-array-counted-by-3.c: New test.
---
 .../gcc.dg/flex-array-counted-by-2.c          |  74 +++++++
 .../gcc.dg/flex-array-counted-by-3.c          | 197 ++++++++++++++++++
 gcc/tree-object-size.cc                       |  37 +++-
 gcc/tree.cc                                   |  95 ++++++++-
 gcc/tree.h                                    |  10 +
 5 files changed, 405 insertions(+), 8 deletions(-)
 create mode 100644 gcc/testsuite/gcc.dg/flex-array-counted-by-2.c
 create mode 100644 gcc/testsuite/gcc.dg/flex-array-counted-by-3.c

diff --git a/gcc/testsuite/gcc.dg/flex-array-counted-by-2.c b/gcc/testsuite/gcc.dg/flex-array-counted-by-2.c
new file mode 100644
index 000000000000..ec580c1f1f01
--- /dev/null
+++ b/gcc/testsuite/gcc.dg/flex-array-counted-by-2.c
@@ -0,0 +1,74 @@
+/* test the attribute counted_by and its usage in
+ * __builtin_dynamic_object_size.  */ 
+/* { dg-do run } */
+/* { dg-options "-O2" } */
+
+#include "builtin-object-size-common.h"
+
+#define expect(p, _v) do { \
+    size_t v = _v; \
+    if (p == v) \
+	__builtin_printf ("ok:  %s == %zd\n", #p, p); \
+    else \
+	{  \
+	  __builtin_printf ("WAT: %s == %zd (expected %zd)\n", #p, p, v); \
+	  FAIL (); \
+	} \
+} while (0);
+
+struct flex {
+  int b;
+  int c[];
+} *array_flex;
+
+struct annotated {
+  int b;
+  int c[] __attribute__ ((counted_by (b)));
+} *array_annotated;
+
+struct nested_annotated {
+  struct {
+    union {
+      int b;
+      float f;	
+    };
+    int n;
+  };
+  int c[] __attribute__ ((counted_by (b)));
+} *array_nested_annotated;
+
+void __attribute__((__noinline__)) setup (int normal_count, int attr_count)
+{
+  array_flex
+    = (struct flex *)malloc (sizeof (struct flex)
+			     + normal_count *  sizeof (int));
+  array_flex->b = normal_count;
+
+  array_annotated
+    = (struct annotated *)malloc (sizeof (struct annotated)
+				  + attr_count *  sizeof (int));
+  array_annotated->b = attr_count;
+
+  array_nested_annotated
+    = (struct nested_annotated *)malloc (sizeof (struct nested_annotated)
+					 + attr_count *  sizeof (int));
+  array_nested_annotated->b = attr_count;
+
+  return;
+}
+
+void __attribute__((__noinline__)) test ()
+{
+    expect(__builtin_dynamic_object_size(array_flex->c, 1), -1);
+    expect(__builtin_dynamic_object_size(array_annotated->c, 1),
+	   array_annotated->b * sizeof (int));
+    expect(__builtin_dynamic_object_size(array_nested_annotated->c, 1),
+	   array_nested_annotated->b * sizeof (int));
+}
+
+int main(int argc, char *argv[])
+{
+  setup (10,10);   
+  test ();
+  DONE ();
+}
diff --git a/gcc/testsuite/gcc.dg/flex-array-counted-by-3.c b/gcc/testsuite/gcc.dg/flex-array-counted-by-3.c
new file mode 100644
index 000000000000..22ef2af31c20
--- /dev/null
+++ b/gcc/testsuite/gcc.dg/flex-array-counted-by-3.c
@@ -0,0 +1,197 @@
+/* test the attribute counted_by and its usage in
+__builtin_dynamic_object_size: what's the correct behavior when the allocaiton
+size mismatched with the value of counted_by attribute?  */
+/* { dg-do run } */
+/* { dg-options "-O -fstrict-flex-arrays=3" } */
+
+#include "builtin-object-size-common.h"
+
+struct annotated {
+  size_t foo;
+  int array[] __attribute__((counted_by (foo)));
+};
+
+#define expect(p, _v) do { \
+    size_t v = _v; \
+    if (p == v) \
+        __builtin_printf ("ok:  %s == %zd\n", #p, p); \
+    else \
+        {  \
+          __builtin_printf ("WAT: %s == %zd (expected %zd)\n", #p, p, v); \
+	  FAIL (); \
+        } \
+} while (0);
+
+#define noinline __attribute__((__noinline__))
+#define SIZE_BUMP 5
+
+/* In general, Due to type casting, the type for the pointee of a pointer
+   does not say anything about the object it points to,
+   So, __builtin_object_size can not directly use the type of the pointee
+   to decide the size of the object the pointer points to.
+
+   there are only two reliable ways:
+   A. observed allocations  (call to the allocation functions in the routine)
+   B. observed accesses     (read or write access to the location of the
+                             pointer points to)
+
+   that provide information about the type/existence of an object at
+   the corresponding address.
+
+   for A, we use the "alloc_size" attribute for the corresponding allocation
+   functions to determine the object size;
+
+   For B, we use the SIZE info of the TYPE attached to the corresponding access.
+   (We treat counted_by attribute as a complement to the SIZE info of the TYPE
+    for FMA)
+
+   The only other way in C which ensures that a pointer actually points
+   to an object of the correct type is 'static':
+
+   void foo(struct P *p[static 1]);
+
+   See https://gcc.gnu.org/pipermail/gcc-patches/2023-July/624814.html
+   for more details.  */
+
+/* in the following function, malloc allocated more space than the value
+   of counted_by attribute.  Then what's the correct behavior we expect
+   the __builtin_dynamic_object_size should have for each of the cases?  */
+
+static struct annotated * noinline alloc_buf_more (int index)
+{
+  struct annotated *p;
+  p = malloc(sizeof (*p) + (index + SIZE_BUMP) * sizeof (int));
+  p->foo = index;
+
+  /*when checking the observed access p->array, we have info on both
+    observered allocation and observed access,
+    A. from observed allocation: (index + SIZE_BUMP) * sizeof (int)
+    B. from observed access: p->foo * sizeof (int)
+
+    in the above, p->foo = index.
+   */
+
+  /* for size in the whole object: always uses A.  */
+  /* for size in the sub-object: chose the smaller of A and B.
+   * Please see https://gcc.gnu.org/pipermail/gcc-patches/2023-July/625891.html
+   * for details on why.  */
+
+  /* for MAXIMUM size in the whole object: use the allocation size
+     for the whole object.  */
+  expect(__builtin_dynamic_object_size(p->array, 0),
+	 (index + SIZE_BUMP) * sizeof(int));
+
+  /* for MAXIMUM size in the sub-object. use the smaller of A and B.  */
+  expect(__builtin_dynamic_object_size(p->array, 1), (p->foo) * sizeof(int));
+
+  /* for MINIMUM size in the whole object: use the allocation size
+     for the whole object.  */
+  expect(__builtin_dynamic_object_size(p->array, 2),
+	 (index + SIZE_BUMP) * sizeof(int));
+
+  /* for MINIMUM size in the sub-object: use the smaller of A and B.  */
+  expect(__builtin_dynamic_object_size(p->array, 3), p->foo * sizeof(int));
+
+  /*when checking the pointer p, we only have info on the observed allocation.
+    So, the object size info can only been obtained from the call to malloc.
+    for both MAXIMUM and MINIMUM: A = (index + SIZE_BUMP) * sizeof (int)  */
+  expect(__builtin_dynamic_object_size(p, 1),
+	 sizeof (*p) + (index + SIZE_BUMP) * sizeof(int));
+  expect(__builtin_dynamic_object_size(p, 0),
+	 sizeof (*p) + (index + SIZE_BUMP) * sizeof(int));
+  expect(__builtin_dynamic_object_size(p, 3),
+	 sizeof (*p) + (index + SIZE_BUMP) * sizeof(int));
+  expect(__builtin_dynamic_object_size(p, 2),
+	 sizeof (*p) + (index + SIZE_BUMP) * sizeof(int));
+  return p;
+}
+
+/* in the following function, malloc allocated less space than the value
+   of counted_by attribute.  Then what's the correct behavior we expect
+   the __builtin_dynamic_object_size should have for each of the cases?
+   NOTE: this is an user error, GCC should issue warnings for such case.
+   this is a seperate issue we should address later.  */
+
+static struct annotated * noinline alloc_buf_less (int index)
+{
+  struct annotated *p;
+  p = malloc(sizeof (*p) + (index) * sizeof (int));
+  p->foo = index + SIZE_BUMP;
+
+  /*when checking the observed access p->array, we have info on both
+    observered allocation and observed access,
+    A. from observed allocation: (index) * sizeof (int)
+    B. from observed access: p->foo * sizeof (int)
+
+    in the above, p->foo = index + SIZE_BUMP.
+   */
+
+  /* for size in the whole object: always uses A.  */
+  /* for size in the sub-object: chose the smaller of A and B.
+   * Please see https://gcc.gnu.org/pipermail/gcc-patches/2023-July/625891.html
+   * for details on why.  */
+
+  /* for MAXIMUM size in the whole object: use the allocation size 
+     for the whole object.  */
+  expect(__builtin_dynamic_object_size(p->array, 0), (index) * sizeof(int));
+
+  /* for MAXIMUM size in the sub-object. use the smaller of A and B.  */ 
+  expect(__builtin_dynamic_object_size(p->array, 1), (index) * sizeof(int));
+
+  /* for MINIMUM size in the whole object: use the allocation size
+     for the whole object.  */
+  expect(__builtin_dynamic_object_size(p->array, 2), (index) * sizeof(int));
+
+  /* for MINIMUM size in the sub-object: use the smaller of A and B.  */
+  expect(__builtin_dynamic_object_size(p->array, 3), (index) * sizeof(int));
+
+  /*when checking the pointer p, we only have info on the observed allocation.
+    So, the object size info can only been obtained from the call to malloc.
+    for both MAXIMUM and MINIMUM: A = (index + SIZE_BUMP) * sizeof (int)  */
+  expect(__builtin_dynamic_object_size(p, 1),
+	 sizeof (*p) + (index) * sizeof(int));
+  expect(__builtin_dynamic_object_size(p, 0),
+	 sizeof (*p) + (index) * sizeof(int));
+  expect(__builtin_dynamic_object_size(p, 3),
+	 sizeof (*p) + (index) * sizeof(int));
+  expect(__builtin_dynamic_object_size(p, 2),
+	 sizeof (*p) + (index) * sizeof(int));
+  return p;
+}
+
+int main ()
+{
+  struct annotated *p, *q;
+  p = alloc_buf_more (10);
+  q = alloc_buf_less (10);
+
+  /*when checking the observed access p->array, we only have info on the
+    observed access, i.e, the TYPE_SIZE info from the access. We don't have
+    info on the whole object.  */
+  expect(__builtin_dynamic_object_size(p->array, 1), p->foo * sizeof(int));
+  expect(__builtin_dynamic_object_size(p->array, 0), -1);
+  expect(__builtin_dynamic_object_size(p->array, 3), p->foo * sizeof(int));
+  expect(__builtin_dynamic_object_size(p->array, 2), 0);
+  /*when checking the pointer p, we have no observed allocation nor observed
+    access, therefore, we cannot determine the size info here.  */
+  expect(__builtin_dynamic_object_size(p, 1), -1);
+  expect(__builtin_dynamic_object_size(p, 0), -1);
+  expect(__builtin_dynamic_object_size(p, 3), 0);
+  expect(__builtin_dynamic_object_size(p, 2), 0);
+
+  /*when checking the observed access p->array, we only have info on the
+    observed access, i.e, the TYPE_SIZE info from the access. We don't have
+    info on the whole object.  */
+  expect(__builtin_dynamic_object_size(q->array, 1), q->foo * sizeof(int));
+  expect(__builtin_dynamic_object_size(q->array, 0), -1);
+  expect(__builtin_dynamic_object_size(q->array, 3), q->foo * sizeof(int));
+  expect(__builtin_dynamic_object_size(q->array, 2), 0);
+  /*when checking the pointer p, we have no observed allocation nor observed
+    access, therefore, we cannot determine the size info here.  */
+  expect(__builtin_dynamic_object_size(q, 1), -1);
+  expect(__builtin_dynamic_object_size(q, 0), -1);
+  expect(__builtin_dynamic_object_size(q, 3), 0);
+  expect(__builtin_dynamic_object_size(q, 2), 0);
+
+  DONE ();
+}
diff --git a/gcc/tree-object-size.cc b/gcc/tree-object-size.cc
index a62af0500563..cf7843c5684b 100644
--- a/gcc/tree-object-size.cc
+++ b/gcc/tree-object-size.cc
@@ -585,6 +585,7 @@ addr_object_size (struct object_size_info *osi, const_tree ptr,
   if (pt_var != TREE_OPERAND (ptr, 0))
     {
       tree var;
+      tree counted_by_ref = NULL_TREE;
 
       if (object_size_type & OST_SUBOBJECT)
 	{
@@ -600,11 +601,12 @@ addr_object_size (struct object_size_info *osi, const_tree ptr,
 	    var = TREE_OPERAND (var, 0);
 	  if (var != pt_var && TREE_CODE (var) == ARRAY_REF)
 	    var = TREE_OPERAND (var, 0);
-	  if (! TYPE_SIZE_UNIT (TREE_TYPE (var))
+	  if (! component_ref_has_counted_by_p (var)
+	     && ((! TYPE_SIZE_UNIT (TREE_TYPE (var))
 	      || ! tree_fits_uhwi_p (TYPE_SIZE_UNIT (TREE_TYPE (var)))
 	      || (pt_var_size && TREE_CODE (pt_var_size) == INTEGER_CST
 		  && tree_int_cst_lt (pt_var_size,
-				      TYPE_SIZE_UNIT (TREE_TYPE (var)))))
+				      TYPE_SIZE_UNIT (TREE_TYPE (var)))))))
 	    var = pt_var;
 	  else if (var != pt_var && TREE_CODE (pt_var) == MEM_REF)
 	    {
@@ -612,6 +614,7 @@ addr_object_size (struct object_size_info *osi, const_tree ptr,
 	      /* For &X->fld, compute object size if fld isn't a flexible array
 		 member.  */
 	      bool is_flexible_array_mem_ref = false;
+
 	      while (v && v != pt_var)
 		switch (TREE_CODE (v))
 		  {
@@ -660,6 +663,8 @@ addr_object_size (struct object_size_info *osi, const_tree ptr,
 		    /* Now the ref is to an array type.  */
 		    gcc_assert (TREE_CODE (TREE_TYPE (v)) == ARRAY_TYPE);
 		    is_flexible_array_mem_ref = array_ref_flexible_size_p (v);
+		    counted_by_ref = component_ref_get_counted_by (v);
+
 		    while (v != pt_var && TREE_CODE (v) == COMPONENT_REF)
 		      if (TREE_CODE (TREE_TYPE (TREE_OPERAND (v, 0)))
 			  != UNION_TYPE
@@ -673,8 +678,11 @@ addr_object_size (struct object_size_info *osi, const_tree ptr,
 			   == RECORD_TYPE)
 		      {
 			/* compute object size only if v is not a
-			   flexible array member.  */
-			if (!is_flexible_array_mem_ref)
+			   flexible array member or the flexible array member
+			   has a known element count indicated by the user
+			   through attribute counted_by.  */
+			if (!is_flexible_array_mem_ref
+			    || counted_by_ref)
 			  {
 			    v = NULL_TREE;
 			    break;
@@ -707,9 +715,24 @@ addr_object_size (struct object_size_info *osi, const_tree ptr,
 
       if (var != pt_var)
 	{
-	  var_size = TYPE_SIZE_UNIT (TREE_TYPE (var));
-	  if (!TREE_CONSTANT (var_size))
-	    var_size = get_or_create_ssa_default_def (cfun, var_size);
+	  if (!counted_by_ref)
+	    {
+	      var_size = TYPE_SIZE_UNIT (TREE_TYPE (var));
+	      if (!TREE_CONSTANT (var_size))
+		var_size = get_or_create_ssa_default_def (cfun, var_size);
+	    }
+	  else
+	    {
+	      gcc_assert (TREE_CODE (var) == COMPONENT_REF
+			  && TREE_CODE (TREE_TYPE (var)) == ARRAY_TYPE);
+	      tree element_size = TYPE_SIZE_UNIT (TREE_TYPE (TREE_TYPE (var)));
+	      var_size
+		= size_binop (MULT_EXPR,
+			      fold_convert (sizetype, counted_by_ref),
+			      fold_convert (sizetype, element_size));
+	      if (!todo)
+		todo = TODO_update_ssa_only_virtuals;
+	    }
 	  if (!var_size)
 	    return false;
 	}
diff --git a/gcc/tree.cc b/gcc/tree.cc
index fcd36ae0cd74..c46b73be3906 100644
--- a/gcc/tree.cc
+++ b/gcc/tree.cc
@@ -12745,6 +12745,32 @@ array_ref_element_size (tree exp)
     return SUBSTITUTE_PLACEHOLDER_IN_EXPR (TYPE_SIZE_UNIT (elmt_type), exp);
 }
 
+/*  For a component_ref that has an array type ARRAY_REF, return TRUE when
+    an counted_by attribute attached to the corresponding FIELD_DECL.
+    return FALSE otherwise.  */
+bool
+component_ref_has_counted_by_p (tree array_ref)
+{
+  if (TREE_CODE (array_ref) != COMPONENT_REF)
+    return false;
+
+  if (TREE_CODE (TREE_TYPE (array_ref)) != ARRAY_TYPE)
+    return false;
+
+  tree struct_object = TREE_OPERAND (array_ref, 0);
+  tree struct_type = TREE_TYPE (struct_object);
+
+  if (!RECORD_OR_UNION_TYPE_P (struct_type))
+    return false;
+  tree field_decl = TREE_OPERAND (array_ref, 1);
+  tree attr_counted_by = lookup_attribute ("counted_by",
+					      DECL_ATTRIBUTES (field_decl));
+
+  if (!attr_counted_by)
+    return false;
+  return true;
+}
+
 /* Given a field list, FIELDLIST, of a structure/union, return a TREE_LIST,
    with each TREE_VALUE a FIELD_DECL stepping down the chain to the FIELD
    whose name is FIELDNAME, which is the last TREE_VALUE of the list.
@@ -12771,7 +12797,7 @@ get_named_field (tree fieldlist, const char *fieldname)
 	 fields inside it recursively.  */
       else if (RECORD_OR_UNION_TYPE_P (TREE_TYPE (field)))
 	  if ((named_field = get_named_field (TYPE_FIELDS (TREE_TYPE (field)),
-						  fieldname)) != NULL_TREE)
+					     fieldname)) != NULL_TREE)
 	    {
 	      named_field = tree_cons (NULL_TREE, field, named_field);
 	      break;
@@ -12784,6 +12810,73 @@ get_named_field (tree fieldlist, const char *fieldname)
   return named_field;
 }
 
+/* For a component_ref that has an array type ARRAY_REF, get the object that
+   represents its counted_by per the attribute counted_by attached to
+   the corresponding FIELD_DECL.  return NULL_TREE when cannot find such
+   object.
+   For example, if:
+
+    struct P {
+      int k;
+      int x[] __attribute__ ((counted_by (k)));
+    } *p;
+
+    for the following reference:
+
+    p->x[b]
+
+    the object that represents its element count will be:
+
+    p->k
+
+    So, when component_ref_get_counted_by (p->x[b]) is called, p->k should be
+    returned.
+*/
+
+tree
+component_ref_get_counted_by (tree array_ref)
+{
+  if (! component_ref_has_counted_by_p (array_ref))
+    return NULL_TREE;
+
+  tree struct_object = TREE_OPERAND (array_ref, 0);
+  tree struct_type = TREE_TYPE (struct_object);
+  tree field_decl = TREE_OPERAND (array_ref, 1);
+  tree attr_counted_by = lookup_attribute ("counted_by",
+					   DECL_ATTRIBUTES (field_decl));
+  gcc_assert (attr_counted_by);
+
+  /* If there is an counted_by attribute attached to the field,
+     get the field that maps to the counted_by.  */
+
+  const char *fieldname
+    = IDENTIFIER_POINTER (TREE_VALUE (TREE_VALUE (attr_counted_by)));
+
+  tree counted_by_field = get_named_field (TYPE_FIELDS (struct_type),
+					   fieldname);
+
+  gcc_assert (counted_by_field);
+
+  /* generate the tree node that represent the counted_by of this array
+     ref.  This is a (possible nested) COMPONENT_REF to the counted_by_field
+     of the containing structure.  */
+
+  tree counted_by_ref = NULL_TREE;
+  tree object = struct_object;
+  do
+    {
+      tree field = TREE_VALUE (counted_by_field);
+
+      counted_by_ref = build3 (COMPONENT_REF,
+				  TREE_TYPE (field),
+				  object, field,
+				  NULL_TREE);
+      object = counted_by_ref;
+      counted_by_field = TREE_CHAIN (counted_by_field);
+    }
+  while (counted_by_field);
+  return counted_by_ref;
+}
 
 /* Return a tree representing the lower bound of the array mentioned in
    EXP, an ARRAY_REF or an ARRAY_RANGE_REF.  */
diff --git a/gcc/tree.h b/gcc/tree.h
index 4859becaa1e7..07eed7219835 100644
--- a/gcc/tree.h
+++ b/gcc/tree.h
@@ -5619,11 +5619,21 @@ extern tree get_base_address (tree t);
    of EXP, an ARRAY_REF or an ARRAY_RANGE_REF.  */
 extern tree array_ref_element_size (tree);
 
+/* Give a component_ref that has an array type, return true when an
+   attribute counted_by attached to the corresponding FIELD_DECL.  */
+extern bool component_ref_has_counted_by_p (tree);
+
 /* Given a field list, FIELDLIST, of a structure/union, return the FIELD whose
    name is FIELDNAME, return NULL_TREE if such field is not found.
    searching nested anonymous structure/union recursively.  */
 extern tree get_named_field (tree, const char *);
 
+/* Give a component_ref that has an array type, return the object that
+   represents its counted_by per the attribute counted_by attached to
+   the corresponding FIELD_DECL.  return NULL_TREE when cannot find such
+   object.  */
+extern tree component_ref_get_counted_by (tree);
+
 /* Return a typenode for the "standard" C type with a given name.  */
 extern tree get_typenode_from_name (const char *);
 
-- 
2.31.1