Actual source code: pheap.c


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

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

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

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

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

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

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

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

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

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

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

 84:   PetscFunctionBegin;
 85:   if (1 < h->end && h->end < ARITY) h->end = ARITY;
 86:   loc = h->end++;
 87:   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));
 88:   h->base[loc].id    = id;
 89:   h->base[loc].value = val;

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

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

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

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

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

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

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

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

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

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

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

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

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

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