1 2 #include <petsc/private/petscimpl.h> 3 #include <petscviewer.h> 4 5 typedef struct { 6 PetscInt id; 7 PetscInt value; 8 } HeapNode; 9 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 }; 16 17 /* 18 * The arity of the heap can be changed via the parameter B below. Consider the B=2 (arity=4 case below) 19 * 20 * [00 (sentinel); 01 (min node); 10 (unused); 11 (unused); 0100 (first child); 0101; 0110; 0111; ...] 21 * 22 * Slots 10 and 11 are referred to as the "hole" below in the implementation. 23 */ 24 25 #define B 1 /* log2(ARITY) */ 26 #define ARITY (1 << B) /* tree branching factor */ 27 static inline PetscInt Parent(PetscInt loc) { 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) 34 35 static inline void Swap(PetscHeap h, PetscInt loc, PetscInt loc2) { 36 PetscInt id, val; 37 id = Id(h, loc); 38 val = Value(h, loc); 39 h->base[loc].id = Id(h, loc2); 40 h->base[loc].value = Value(h, loc2); 41 h->base[loc2].id = id; 42 h->base[loc2].value = val; 43 } 44 static inline PetscInt MinChild(PetscHeap h, PetscInt loc) { 45 PetscInt min, chld, left, right; 46 left = loc << B; 47 right = PetscMin(left + ARITY - 1, h->end - 1); 48 chld = 0; 49 min = PETSC_MAX_INT; 50 for (; left <= right; left++) { 51 PetscInt val = Value(h, left); 52 if (val < min) { 53 min = val; 54 chld = left; 55 } 56 } 57 return chld; 58 } 59 60 PetscErrorCode PetscHeapCreate(PetscInt maxsize, PetscHeap *heap) { 61 PetscHeap h; 62 63 PetscFunctionBegin; 64 *heap = NULL; 65 PetscCall(PetscMalloc1(1, &h)); 66 h->end = 1; 67 h->alloc = maxsize + ARITY; /* We waste all but one slot (loc=1) in the first ARITY slots */ 68 h->stash = h->alloc; 69 PetscCall(PetscCalloc1(h->alloc, &h->base)); 70 h->base[0].id = -1; 71 h->base[0].value = PETSC_MIN_INT; 72 *heap = h; 73 PetscFunctionReturn(0); 74 } 75 76 PetscErrorCode PetscHeapAdd(PetscHeap h, PetscInt id, PetscInt val) { 77 PetscInt loc, par; 78 79 PetscFunctionBegin; 80 if (1 < h->end && h->end < ARITY) h->end = ARITY; 81 loc = h->end++; 82 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)); 83 h->base[loc].id = id; 84 h->base[loc].value = val; 85 86 /* move up until heap condition is satisfied */ 87 while ((void)(par = Parent(loc)), Value(h, par) > val) { 88 Swap(h, loc, par); 89 loc = par; 90 } 91 PetscFunctionReturn(0); 92 } 93 94 PetscErrorCode PetscHeapPop(PetscHeap h, PetscInt *id, PetscInt *val) { 95 PetscInt loc, chld; 96 97 PetscFunctionBegin; 98 if (h->end == 1) { 99 *id = h->base[0].id; 100 *val = h->base[0].value; 101 PetscFunctionReturn(0); 102 } 103 104 *id = h->base[1].id; 105 *val = h->base[1].value; 106 107 /* rotate last entry into first position */ 108 loc = --h->end; 109 if (h->end == ARITY) h->end = 2; /* Skip over hole */ 110 h->base[1].id = Id(h, loc); 111 h->base[1].value = Value(h, loc); 112 113 /* move down until min heap condition is satisfied */ 114 loc = 1; 115 while ((chld = MinChild(h, loc)) && Value(h, loc) > Value(h, chld)) { 116 Swap(h, loc, chld); 117 loc = chld; 118 } 119 PetscFunctionReturn(0); 120 } 121 122 PetscErrorCode PetscHeapPeek(PetscHeap h, PetscInt *id, PetscInt *val) { 123 PetscFunctionBegin; 124 if (h->end == 1) { 125 *id = h->base[0].id; 126 *val = h->base[0].value; 127 PetscFunctionReturn(0); 128 } 129 130 *id = h->base[1].id; 131 *val = h->base[1].value; 132 PetscFunctionReturn(0); 133 } 134 135 PetscErrorCode PetscHeapStash(PetscHeap h, PetscInt id, PetscInt val) { 136 PetscInt loc; 137 138 PetscFunctionBegin; 139 loc = --h->stash; 140 h->base[loc].id = id; 141 h->base[loc].value = val; 142 PetscFunctionReturn(0); 143 } 144 145 PetscErrorCode PetscHeapUnstash(PetscHeap h) { 146 PetscFunctionBegin; 147 while (h->stash < h->alloc) { 148 PetscInt id = Id(h, h->stash), value = Value(h, h->stash); 149 h->stash++; 150 PetscCall(PetscHeapAdd(h, id, value)); 151 } 152 PetscFunctionReturn(0); 153 } 154 155 PetscErrorCode PetscHeapDestroy(PetscHeap *heap) { 156 PetscFunctionBegin; 157 PetscCall(PetscFree((*heap)->base)); 158 PetscCall(PetscFree(*heap)); 159 PetscFunctionReturn(0); 160 } 161 162 PetscErrorCode PetscHeapView(PetscHeap h, PetscViewer viewer) { 163 PetscBool iascii; 164 165 PetscFunctionBegin; 166 if (!viewer) PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_SELF, &viewer)); 167 PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2); 168 PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii)); 169 if (iascii) { 170 PetscCall(PetscViewerASCIIPrintf(viewer, "Heap size %" PetscInt_FMT " with %" PetscInt_FMT " stashed\n", h->end - 1, h->alloc - h->stash)); 171 PetscCall(PetscViewerASCIIPrintf(viewer, "Heap in (id,value) pairs\n")); 172 PetscCall(PetscIntView(2 * (h->end - 1), (const PetscInt *)(h->base + 1), viewer)); 173 PetscCall(PetscViewerASCIIPrintf(viewer, "Stash in (id,value) pairs\n")); 174 PetscCall(PetscIntView(2 * (h->alloc - h->stash), (const PetscInt *)(h->base + h->stash), viewer)); 175 } 176 PetscFunctionReturn(0); 177 } 178