xref: /petsc/src/mat/utils/pheap.c (revision f97672e55eacc8688507b9471cd7ec2664d7f203)
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 {
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)
35 
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 }
62 
63 PetscErrorCode PetscHeapCreate(PetscInt maxsize,PetscHeap *heap)
64 {
65   PetscHeap      h;
66 
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(0);
78 }
79 
80 PetscErrorCode PetscHeapAdd(PetscHeap h,PetscInt id,PetscInt val)
81 {
82   PetscInt loc,par;
83 
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;
90 
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(0);
97 }
98 
99 PetscErrorCode PetscHeapPop(PetscHeap h,PetscInt *id,PetscInt *val)
100 {
101   PetscInt loc,chld;
102 
103   PetscFunctionBegin;
104   if (h->end == 1) {
105     *id  = h->base[0].id;
106     *val = h->base[0].value;
107     PetscFunctionReturn(0);
108   }
109 
110   *id  = h->base[1].id;
111   *val = h->base[1].value;
112 
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);
118 
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(0);
126 }
127 
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(0);
135   }
136 
137   *id  = h->base[1].id;
138   *val = h->base[1].value;
139   PetscFunctionReturn(0);
140 }
141 
142 PetscErrorCode PetscHeapStash(PetscHeap h,PetscInt id,PetscInt val)
143 {
144   PetscInt loc;
145 
146   PetscFunctionBegin;
147   loc                = --h->stash;
148   h->base[loc].id    = id;
149   h->base[loc].value = val;
150   PetscFunctionReturn(0);
151 }
152 
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(0);
162 }
163 
164 PetscErrorCode PetscHeapDestroy(PetscHeap *heap)
165 {
166   PetscFunctionBegin;
167   PetscCall(PetscFree((*heap)->base));
168   PetscCall(PetscFree(*heap));
169   PetscFunctionReturn(0);
170 }
171 
172 PetscErrorCode PetscHeapView(PetscHeap h,PetscViewer viewer)
173 {
174   PetscBool      iascii;
175 
176   PetscFunctionBegin;
177   if (!viewer) {
178     PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_SELF,&viewer));
179   }
180   PetscValidHeaderSpecific(viewer,PETSC_VIEWER_CLASSID,2);
181   PetscCall(PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii));
182   if (iascii) {
183     PetscCall(PetscViewerASCIIPrintf(viewer,"Heap size %" PetscInt_FMT " with %" PetscInt_FMT " stashed\n",h->end-1,h->alloc-h->stash));
184     PetscCall(PetscViewerASCIIPrintf(viewer,"Heap in (id,value) pairs\n"));
185     PetscCall(PetscIntView(2*(h->end-1),(const PetscInt*)(h->base+1),viewer));
186     PetscCall(PetscViewerASCIIPrintf(viewer,"Stash in (id,value) pairs\n"));
187     PetscCall(PetscIntView(2*(h->alloc-h->stash),(const PetscInt*)(h->base+h->stash),viewer));
188   }
189   PetscFunctionReturn(0);
190 }
191