xref: /petsc/src/mat/utils/pheap.c (revision bcda9346efad4e5ba2d553af84eb238771ba1e25)
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 */
Parent(PetscInt loc)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 
Swap(PetscHeap h,PetscInt loc,PetscInt loc2)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 }
MinChild(PetscHeap h,PetscInt loc)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 
PetscHeapCreate(PetscInt maxsize,PetscHeap * heap)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 
PetscHeapAdd(PetscHeap h,PetscInt id,PetscInt val)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 
PetscHeapPop(PetscHeap h,PetscInt * id,PetscInt * val)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 
PetscHeapPeek(PetscHeap h,PetscInt * id,PetscInt * val)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 
PetscHeapStash(PetscHeap h,PetscInt id,PetscInt val)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 
PetscHeapUnstash(PetscHeap h)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 
PetscHeapDestroy(PetscHeap * heap)163 PetscErrorCode PetscHeapDestroy(PetscHeap *heap)
164 {
165   PetscFunctionBegin;
166   PetscCall(PetscFree((*heap)->base));
167   PetscCall(PetscFree(*heap));
168   PetscFunctionReturn(PETSC_SUCCESS);
169 }
170 
PetscHeapView(PetscHeap h,PetscViewer viewer)171 PetscErrorCode PetscHeapView(PetscHeap h, PetscViewer viewer)
172 {
173   PetscBool isascii;
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, &isascii));
179   if (isascii) {
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