mesa/src/util/sparse_array.c

/*
 * Copyright © 2019 Intel Corporation
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice (including the next
 * paragraph) shall be included in all copies or substantial portions of the
 * Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 * IN THE SOFTWARE.
 */

#include "sparse_array.h"

struct util_sparse_array_node {
   uint32_t level;
   uint32_t _pad;
   uint64_t max_idx;
};

void
util_sparse_array_init(struct util_sparse_array *arr,
                       size_t elem_size, size_t node_size)
{
   memset(arr, 0, sizeof(*arr));
   arr->elem_size = elem_size;
   arr->node_size_log2 = util_logbase2_64(node_size);
   assert(node_size >= 2 && node_size == (1ull << arr->node_size_log2));
}

static inline void *
_util_sparse_array_node_data(struct util_sparse_array_node *node)
{
   return node + 1;
}

static inline void
_util_sparse_array_node_finish(struct util_sparse_array *arr,
                               struct util_sparse_array_node *node)
{
   if (node->level > 0) {
      struct util_sparse_array_node **children =
         _util_sparse_array_node_data(node);
      size_t node_size = 1ull << arr->node_size_log2;
      for (size_t i = 0; i < node_size; i++) {
         if (children[i] != NULL)
            _util_sparse_array_node_finish(arr, children[i]);
      }
   }

   free(node);
}

void
util_sparse_array_finish(struct util_sparse_array *arr)
{
   if (arr->root)
      _util_sparse_array_node_finish(arr, arr->root);
}

static inline struct util_sparse_array_node *
_util_sparse_array_alloc_node(struct util_sparse_array *arr,
                              unsigned level)
{
   size_t size = sizeof(struct util_sparse_array_node);
   if (level == 0) {
      size += arr->elem_size << arr->node_size_log2;
   } else {
      size += sizeof(struct util_sparse_array_node *) << arr->node_size_log2;
   }

   struct util_sparse_array_node *node = calloc(1, size);
   node->level = level;

   return node;
}

static inline struct util_sparse_array_node *
_util_sparse_array_set_or_free_node(struct util_sparse_array_node **node_ptr,
                                    struct util_sparse_array_node *cmp_node,
                                    struct util_sparse_array_node *node)
{
   struct util_sparse_array_node *prev_node =
      p_atomic_cmpxchg(node_ptr, cmp_node, node);

   if (prev_node != cmp_node) {
      /* We lost the race.  Free this one and return the one that was already
       * allocated.
       */
      free(node);
      return prev_node;
   } else {
      return node;
   }
}

void *
util_sparse_array_get(struct util_sparse_array *arr, uint64_t idx)
{
   struct util_sparse_array_node *root = p_atomic_read(&arr->root);
   if (unlikely(root == NULL)) {
      unsigned root_level = 0;
      uint64_t idx_iter = idx >> arr->node_size_log2;
      while (idx_iter) {
         idx_iter >>= arr->node_size_log2;
         root_level++;
      }
      struct util_sparse_array_node *new_root =
         _util_sparse_array_alloc_node(arr, root_level);
      root = _util_sparse_array_set_or_free_node(&arr->root, NULL, new_root);
   }

   while (1) {
      uint64_t root_idx = idx >> (root->level * arr->node_size_log2);
      if (likely(root_idx < (1ull << arr->node_size_log2)))
         break;

      /* In this case, we have a root but its level is low enough that the
       * requested index is out-of-bounds.
       */
      struct util_sparse_array_node *new_root =
         _util_sparse_array_alloc_node(arr, root->level + 1);

      struct util_sparse_array_node **new_root_children =
         _util_sparse_array_node_data(new_root);
      new_root_children[0] = root;

      /* We only add one at a time instead of the whole tree because it's
       * easier to ensure correctness of both the tree building and the
       * clean-up path.  Because we're only adding one node we never have to
       * worry about trying to free multiple things without freeing the old
       * things.
       */
      root = _util_sparse_array_set_or_free_node(&arr->root, root, new_root);
   }

   struct util_sparse_array_node *node = root;
   while (node->level > 0) {
      uint64_t child_idx = (idx >> (node->level * arr->node_size_log2)) &
                           ((1ull << arr->node_size_log2) - 1);

      struct util_sparse_array_node **children =
         _util_sparse_array_node_data(node);
      struct util_sparse_array_node *child =
         p_atomic_read(&children[child_idx]);

      if (unlikely(child == NULL)) {
         child = _util_sparse_array_alloc_node(arr, node->level - 1);
         child = _util_sparse_array_set_or_free_node(&children[child_idx],
                                                     NULL, child);
      }

      node = child;
   }

   uint64_t elem_idx = idx & ((1ull << arr->node_size_log2) - 1);
   return (void *)((char *)_util_sparse_array_node_data(node) +
                   (elem_idx * arr->elem_size));
}

static void
validate_node_level(struct util_sparse_array *arr,
                    struct util_sparse_array_node *node,
                    unsigned level)
{
   assert(node->level == level);

   if (node->level > 0) {
      struct util_sparse_array_node **children =
         _util_sparse_array_node_data(node);
      size_t node_size = 1ull << arr->node_size_log2;
      for (size_t i = 0; i < node_size; i++) {
         if (children[i] != NULL)
            validate_node_level(arr, children[i], level - 1);
      }
   }
}

void
util_sparse_array_validate(struct util_sparse_array *arr)
{
   validate_node_level(arr, arr->root, arr->root->level);
}