1 #include <petsc/private/petscimpl.h> 2 #include <petscviewer.h> 3 4 typedef struct { 5 PetscInt id; 6 PetscInt value; 7 } HeapNode; 8 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 }; 15 16 /* 17 The arity of the heap can be changed via the parameter B below. Consider the B=2 (arity=4 case below) 18 19 [00 (sentinel); 01 (min node); 10 (unused); 11 (unused); 0100 (first child); 0101; 0110; 0111; ...] 20 21 Slots 10 and 11 are referred to as the "hole" below in the implementation. 22 */ 23 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) 34 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 } 61 62 PetscErrorCode PetscHeapCreate(PetscInt maxsize, PetscHeap *heap) 63 { 64 PetscHeap h; 65 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 } 78 79 PetscErrorCode PetscHeapAdd(PetscHeap h, PetscInt id, PetscInt val) 80 { 81 PetscInt loc, par; 82 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; 89 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 } 97 98 PetscErrorCode PetscHeapPop(PetscHeap h, PetscInt *id, PetscInt *val) 99 { 100 PetscInt loc, chld; 101 102 PetscFunctionBegin; 103 if (h->end == 1) { 104 *id = h->base[0].id; 105 *val = h->base[0].value; 106 PetscFunctionReturn(PETSC_SUCCESS); 107 } 108 109 *id = h->base[1].id; 110 *val = h->base[1].value; 111 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); 117 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 } 126 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 } 135 136 *id = h->base[1].id; 137 *val = h->base[1].value; 138 PetscFunctionReturn(PETSC_SUCCESS); 139 } 140 141 PetscErrorCode PetscHeapStash(PetscHeap h, PetscInt id, PetscInt val) 142 { 143 PetscInt loc; 144 145 PetscFunctionBegin; 146 loc = --h->stash; 147 h->base[loc].id = id; 148 h->base[loc].value = val; 149 PetscFunctionReturn(PETSC_SUCCESS); 150 } 151 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 } 162 163 PetscErrorCode PetscHeapDestroy(PetscHeap *heap) 164 { 165 PetscFunctionBegin; 166 PetscCall(PetscFree((*heap)->base)); 167 PetscCall(PetscFree(*heap)); 168 PetscFunctionReturn(PETSC_SUCCESS); 169 } 170 171 PetscErrorCode PetscHeapView(PetscHeap h, PetscViewer viewer) 172 { 173 PetscBool iascii; 174 175 PetscFunctionBegin; 176 if (!viewer) PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_SELF, &viewer)); 177 PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2); 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 } 188