Date: 2025-03-29T05:58:39+00:00

Actual source code: pheap.c

      1: #include <petsc/private/petscimpl.h>
      2: #include <petscviewer.h>

      4: typedef struct {
      5:   PetscInt id;
      6:   PetscInt value;
      7: } HeapNode;

      9: struct _n_PetscHeap {
     10:   PetscInt  end;   /* one past the last item */
     11:   PetscInt  alloc; /* length of array */
     12:   PetscInt  stash; /* stash grows down, this points to last item */
     13:   HeapNode *base;
     14: };

     16: /*
     17:   The arity of the heap can be changed via the parameter B below. Consider the B=2 (arity=4 case below)

     19:   [00 (sentinel); 01 (min node); 10 (unused); 11 (unused); 0100 (first child); 0101; 0110; 0111; ...]

     21:   Slots 10 and 11 are referred to as the "hole" below in the implementation.
     22: */

     24: #define B     1        /* log2(ARITY) */
     25: #define ARITY (1 << B) /* tree branching factor */
     26: static inline PetscInt Parent(PetscInt loc)
     27: {
     28:   PetscInt p = loc >> B;
     29:   if (p < ARITY) return (PetscInt)(loc != 1); /* Parent(1) is 0, otherwise fix entries ending up in the hole */
     30:   return p;
     31: }
     32: #define Value(h, loc) ((h)->base[loc].value)
     33: #define Id(h, loc)    ((h)->base[loc].id)

     35: static inline void Swap(PetscHeap h, PetscInt loc, PetscInt loc2)
     36: {
     37:   PetscInt id, val;
     38:   id                  = Id(h, loc);
     39:   val                 = Value(h, loc);
     40:   h->base[loc].id     = Id(h, loc2);
     41:   h->base[loc].value  = Value(h, loc2);
     42:   h->base[loc2].id    = id;
     43:   h->base[loc2].value = val;
     44: }
     45: static inline PetscInt MinChild(PetscHeap h, PetscInt loc)
     46: {
     47:   PetscInt min, chld, left, right;
     48:   left  = loc << B;
     49:   right = PetscMin(left + ARITY - 1, h->end - 1);
     50:   chld  = 0;
     51:   min   = PETSC_INT_MAX;
     52:   for (; left <= right; left++) {
     53:     PetscInt val = Value(h, left);
     54:     if (val < min) {
     55:       min  = val;
     56:       chld = left;
     57:     }
     58:   }
     59:   return chld;
     60: }

     62: PetscErrorCode PetscHeapCreate(PetscInt maxsize, PetscHeap *heap)
     63: {
     64:   PetscHeap h;

     66:   PetscFunctionBegin;
     67:   *heap = NULL;
     68:   PetscCall(PetscMalloc1(1, &h));
     69:   h->end   = 1;
     70:   h->alloc = maxsize + ARITY; /* We waste all but one slot (loc=1) in the first ARITY slots */
     71:   h->stash = h->alloc;
     72:   PetscCall(PetscCalloc1(h->alloc, &h->base));
     73:   h->base[0].id    = -1;
     74:   h->base[0].value = PETSC_INT_MIN;
     75:   *heap            = h;
     76:   PetscFunctionReturn(PETSC_SUCCESS);
     77: }

     79: PetscErrorCode PetscHeapAdd(PetscHeap h, PetscInt id, PetscInt val)
     80: {
     81:   PetscInt loc, par;

     83:   PetscFunctionBegin;
     84:   if (1 < h->end && h->end < ARITY) h->end = ARITY;
     85:   loc = h->end++;
     86:   PetscCheck(h->end <= h->stash, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Addition would exceed allocation %" PetscInt_FMT " (%" PetscInt_FMT " stashed)", h->alloc, h->alloc - h->stash);
     87:   h->base[loc].id    = id;
     88:   h->base[loc].value = val;

     90:   /* move up until heap condition is satisfied */
     91:   while ((void)(par = Parent(loc)), Value(h, par) > val) {
     92:     Swap(h, loc, par);
     93:     loc = par;
     94:   }
     95:   PetscFunctionReturn(PETSC_SUCCESS);
     96: }

     98: PetscErrorCode PetscHeapPop(PetscHeap h, PetscInt *id, PetscInt *val)
     99: {
    100:   PetscInt loc, chld;

    102:   PetscFunctionBegin;
    103:   if (h->end == 1) {
    104:     *id  = h->base[0].id;
    105:     *val = h->base[0].value;
    106:     PetscFunctionReturn(PETSC_SUCCESS);
    107:   }

    109:   *id  = h->base[1].id;
    110:   *val = h->base[1].value;

    112:   /* rotate last entry into first position */
    113:   loc = --h->end;
    114:   if (h->end == ARITY) h->end = 2; /* Skip over hole */
    115:   h->base[1].id    = Id(h, loc);
    116:   h->base[1].value = Value(h, loc);

    118:   /* move down until min heap condition is satisfied */
    119:   loc = 1;
    120:   while ((chld = MinChild(h, loc)) && Value(h, loc) > Value(h, chld)) {
    121:     Swap(h, loc, chld);
    122:     loc = chld;
    123:   }
    124:   PetscFunctionReturn(PETSC_SUCCESS);
    125: }

    127: PetscErrorCode PetscHeapPeek(PetscHeap h, PetscInt *id, PetscInt *val)
    128: {
    129:   PetscFunctionBegin;
    130:   if (h->end == 1) {
    131:     *id  = h->base[0].id;
    132:     *val = h->base[0].value;
    133:     PetscFunctionReturn(PETSC_SUCCESS);
    134:   }

    136:   *id  = h->base[1].id;
    137:   *val = h->base[1].value;
    138:   PetscFunctionReturn(PETSC_SUCCESS);
    139: }

    141: PetscErrorCode PetscHeapStash(PetscHeap h, PetscInt id, PetscInt val)
    142: {
    143:   PetscInt loc;

    145:   PetscFunctionBegin;
    146:   loc                = --h->stash;
    147:   h->base[loc].id    = id;
    148:   h->base[loc].value = val;
    149:   PetscFunctionReturn(PETSC_SUCCESS);
    150: }

    152: PetscErrorCode PetscHeapUnstash(PetscHeap h)
    153: {
    154:   PetscFunctionBegin;
    155:   while (h->stash < h->alloc) {
    156:     PetscInt id = Id(h, h->stash), value = Value(h, h->stash);
    157:     h->stash++;
    158:     PetscCall(PetscHeapAdd(h, id, value));
    159:   }
    160:   PetscFunctionReturn(PETSC_SUCCESS);
    161: }

    163: PetscErrorCode PetscHeapDestroy(PetscHeap *heap)
    164: {
    165:   PetscFunctionBegin;
    166:   PetscCall(PetscFree((*heap)->base));
    167:   PetscCall(PetscFree(*heap));
    168:   PetscFunctionReturn(PETSC_SUCCESS);
    169: }

    171: PetscErrorCode PetscHeapView(PetscHeap h, PetscViewer viewer)
    172: {
    173:   PetscBool iascii;

    175:   PetscFunctionBegin;
    176:   if (!viewer) PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_SELF, &viewer));
    178:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
    179:   if (iascii) {
    180:     PetscCall(PetscViewerASCIIPrintf(viewer, "Heap size %" PetscInt_FMT " with %" PetscInt_FMT " stashed\n", h->end - 1, h->alloc - h->stash));
    181:     PetscCall(PetscViewerASCIIPrintf(viewer, "Heap in (id,value) pairs\n"));
    182:     PetscCall(PetscIntView(2 * (h->end - 1), (const PetscInt *)(h->base + 1), viewer));
    183:     PetscCall(PetscViewerASCIIPrintf(viewer, "Stash in (id,value) pairs\n"));
    184:     PetscCall(PetscIntView(2 * (h->alloc - h->stash), (const PetscInt *)(h->base + h->stash), viewer));
    185:   }
    186:   PetscFunctionReturn(PETSC_SUCCESS);
    187: }