xref: /petsc/src/mat/utils/pheap.c (revision f1580f4e3ce5d5b2393648fd039d0d41b440385d)
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