1!--------------------------------------------------------------------------------------------------!
2! Copyright (C) by the DBCSR developers group - All rights reserved                                !
3! This file is part of the DBCSR library.                                                          !
4!                                                                                                  !
5! For information on the license, see the LICENSE file.                                            !
6! For further information please visit https://dbcsr.cp2k.org                                      !
7! SPDX-License-Identifier: GPL-2.0+                                                                !
8!--------------------------------------------------------------------------------------------------!
9
10MODULE dbcsr_min_heap
11   USE dbcsr_kinds, ONLY: int_4
12#include "base/dbcsr_base_uses.f90"
13
14   IMPLICIT NONE
15   PRIVATE
16
17   PUBLIC :: dbcsr_heap_type, keyt, valt
18   PUBLIC :: dbcsr_heap_pop, dbcsr_heap_reset_node, dbcsr_heap_fill
19   PUBLIC :: dbcsr_heap_new, dbcsr_heap_release
20   PUBLIC :: dbcsr_heap_get_first, dbcsr_heap_reset_first
21
22   ! Sets the types
23   INTEGER, PARAMETER         :: keyt = int_4
24   INTEGER, PARAMETER         :: valt = int_4
25
26   TYPE dbcsr_heap_node
27      INTEGER(KIND=keyt) :: key
28      INTEGER(KIND=valt) :: value
29   END TYPE dbcsr_heap_node
30
31   TYPE dbcsr_heap_node_e
32      TYPE(dbcsr_heap_node) :: node
33   END TYPE dbcsr_heap_node_e
34
35   TYPE dbcsr_heap_type
36      INTEGER :: n
37      INTEGER, DIMENSION(:), POINTER           :: index
38      TYPE(dbcsr_heap_node_e), DIMENSION(:), POINTER :: nodes
39   END TYPE dbcsr_heap_type
40
41CONTAINS
42
43   ! Lookup functions
44
45   ELEMENTAL FUNCTION get_parent(n) RESULT(parent)
46      INTEGER, INTENT(IN)                                :: n
47      INTEGER                                            :: parent
48
49      parent = INT(n/2)
50   END FUNCTION get_parent
51
52   ELEMENTAL FUNCTION get_left_child(n) RESULT(child)
53      INTEGER, INTENT(IN)                                :: n
54      INTEGER                                            :: child
55
56      child = 2*n
57   END FUNCTION get_left_child
58
59   ELEMENTAL FUNCTION get_value(heap, n) RESULT(value)
60      TYPE(dbcsr_heap_type), INTENT(IN)                  :: heap
61      INTEGER, INTENT(IN)                                :: n
62      INTEGER(KIND=valt)                                 :: value
63
64      value = heap%nodes(n)%node%value
65   END FUNCTION get_value
66
67   ! Initialization functions
68
69   SUBROUTINE dbcsr_heap_new(heap, n)
70      TYPE(dbcsr_heap_type), INTENT(OUT)                 :: heap
71      INTEGER, INTENT(IN)                                :: n
72
73      heap%n = n
74      ALLOCATE (heap%index(n))
75      ALLOCATE (heap%nodes(n))
76   END SUBROUTINE dbcsr_heap_new
77
78   SUBROUTINE dbcsr_heap_release(heap)
79      TYPE(dbcsr_heap_type), INTENT(INOUT)               :: heap
80
81      DEALLOCATE (heap%index)
82      DEALLOCATE (heap%nodes)
83      heap%n = 0
84   END SUBROUTINE dbcsr_heap_release
85
86   SUBROUTINE dbcsr_heap_fill(heap, values)
87      !! Fill heap with given values
88      TYPE(dbcsr_heap_type), INTENT(INOUT)               :: heap
89      INTEGER(KIND=valt), DIMENSION(:), INTENT(IN)       :: values
90
91      INTEGER                                            :: first, i, n
92
93      n = SIZE(values)
94      DBCSR_ASSERT(heap%n >= n)
95
96      DO i = 1, n
97         heap%index(i) = i
98         heap%nodes(i)%node%key = i
99         heap%nodes(i)%node%value = values(i)
100      END DO
101      ! Sort from the last full subtree
102      first = get_parent(n)
103      DO i = first, 1, -1
104         CALL bubble_down(heap, i)
105      ENDDO
106
107   END SUBROUTINE dbcsr_heap_fill
108
109   SUBROUTINE dbcsr_heap_get_first(heap, key, value, found)
110      !! Returns the first heap element without removing it.
111      TYPE(dbcsr_heap_type), INTENT(INOUT)               :: heap
112      INTEGER(KIND=keyt), INTENT(OUT)                    :: key
113      INTEGER(KIND=valt), INTENT(OUT)                    :: value
114      LOGICAL, INTENT(OUT)                               :: found
115
116      IF (heap%n .LT. 1) THEN
117         found = .FALSE.
118      ELSE
119         found = .TRUE.
120         key = heap%nodes(1)%node%key
121         value = heap%nodes(1)%node%value
122      ENDIF
123   END SUBROUTINE dbcsr_heap_get_first
124
125   SUBROUTINE dbcsr_heap_pop(heap, key, value, found)
126      !! Returns and removes the first heap element and rebalances
127      !! the heap.
128
129      TYPE(dbcsr_heap_type), INTENT(INOUT)               :: heap
130      INTEGER(KIND=keyt), INTENT(OUT)                    :: key
131      INTEGER(KIND=valt), INTENT(OUT)                    :: value
132      LOGICAL, INTENT(OUT)                               :: found
133
134!
135
136      CALL dbcsr_heap_get_first(heap, key, value, found)
137      IF (found) THEN
138         IF (heap%n .GT. 1) THEN
139            CALL dbcsr_heap_copy_node(heap, 1, heap%n)
140            heap%n = heap%n - 1
141            CALL bubble_down(heap, 1)
142         ELSE
143            heap%n = heap%n - 1
144         ENDIF
145      ENDIF
146   END SUBROUTINE dbcsr_heap_pop
147
148   SUBROUTINE dbcsr_heap_reset_node(heap, key, value)
149      !! Changes the value of the heap element with given key and
150      !! rebalances the heap.
151
152      TYPE(dbcsr_heap_type), INTENT(INOUT)               :: heap
153      INTEGER(KIND=keyt), INTENT(IN)                     :: key
154      INTEGER(KIND=valt), INTENT(IN)                     :: value
155
156      INTEGER                                            :: n, new_pos
157
158      DBCSR_ASSERT(heap%n > 0)
159
160      n = heap%index(key)
161      DBCSR_ASSERT(heap%nodes(n)%node%key == key)
162      heap%nodes(n)%node%value = value
163      CALL bubble_up(heap, n, new_pos)
164      CALL bubble_down(heap, new_pos)
165   END SUBROUTINE dbcsr_heap_reset_node
166
167   SUBROUTINE dbcsr_heap_reset_first(heap, value)
168      !! Changes the value of the minimum heap element and rebalances the heap.
169      TYPE(dbcsr_heap_type), INTENT(INOUT)               :: heap
170      INTEGER(KIND=valt), INTENT(IN)                     :: value
171
172      DBCSR_ASSERT(heap%n > 0)
173      heap%nodes(1)%node%value = value
174      CALL bubble_down(heap, 1)
175   END SUBROUTINE dbcsr_heap_reset_first
176
177   PURE SUBROUTINE dbcsr_heap_swap(heap, e1, e2)
178      TYPE(dbcsr_heap_type), INTENT(INOUT)               :: heap
179      INTEGER, INTENT(IN)                                :: e1, e2
180
181      INTEGER(KIND=keyt)                                 :: key1, key2
182      TYPE(dbcsr_heap_node)                              :: tmp_node
183
184      key1 = heap%nodes(e1)%node%key
185      key2 = heap%nodes(e2)%node%key
186
187      tmp_node = heap%nodes(e1)%node
188      heap%nodes(e1)%node = heap%nodes(e2)%node
189      heap%nodes(e2)%node = tmp_node
190
191      heap%index(key1) = e2
192      heap%index(key2) = e1
193   END SUBROUTINE dbcsr_heap_swap
194
195   PURE SUBROUTINE dbcsr_heap_copy_node(heap, e1, e2)
196      !! Sets node e1 to e2
197      TYPE(dbcsr_heap_type), INTENT(INOUT)               :: heap
198      INTEGER, INTENT(IN)                                :: e1, e2
199
200      INTEGER(KIND=keyt)                                 :: key1, key2
201
202      key1 = heap%nodes(e1)%node%key
203      key2 = heap%nodes(e2)%node%key
204
205      heap%nodes(e1)%node = heap%nodes(e2)%node
206
207      heap%index(key1) = 0
208      heap%index(key2) = e1
209   END SUBROUTINE dbcsr_heap_copy_node
210
211   SUBROUTINE bubble_down(heap, first)
212      !! Balances a heap by bubbling down from the given element.
213      TYPE(dbcsr_heap_type), INTENT(INOUT)               :: heap
214      INTEGER, INTENT(IN)                                :: first
215
216      INTEGER                                            :: e, left_child, right_child, smallest
217      INTEGER(kind=valt)                                 :: left_child_value, min_value, &
218                                                            right_child_value
219      LOGICAL                                            :: all_done
220
221!
222      DBCSR_ASSERT(0 < first .AND. first <= heap%n)
223
224      e = first
225      all_done = .FALSE.
226      ! Check whether we are finished, i.e,. whether the element to
227      ! bubble down is childless.
228      DO WHILE (e .LE. get_parent(heap%n) .AND. .NOT. all_done)
229         ! Determines which node (current, left, or right child) has the
230         ! smallest value.
231         smallest = e
232         min_value = get_value(heap, e)
233         left_child = get_left_child(e)
234         IF (left_child .LE. heap%n) THEN
235            left_child_value = get_value(heap, left_child)
236            IF (left_child_value .LT. min_value) THEN
237               min_value = left_child_value
238               smallest = left_child
239            ENDIF
240         ENDIF
241         right_child = left_child + 1
242         IF (right_child .LE. heap%n) THEN
243            right_child_value = get_value(heap, right_child)
244            IF (right_child_value .LT. min_value) THEN
245               min_value = right_child_value
246               smallest = right_child
247            ENDIF
248         ENDIF
249         !
250         CALL dbcsr_heap_swap(heap, e, smallest)
251         IF (smallest .EQ. e) THEN
252            all_done = .TRUE.
253         ELSE
254            e = smallest
255         ENDIF
256      ENDDO
257   END SUBROUTINE bubble_down
258
259   SUBROUTINE bubble_up(heap, first, new_pos)
260      !! Balances a heap by bubbling up from the given element.
261      TYPE(dbcsr_heap_type), INTENT(INOUT)               :: heap
262      INTEGER, INTENT(IN)                                :: first
263      INTEGER, INTENT(OUT)                               :: new_pos
264
265      INTEGER                                            :: e, parent
266      INTEGER(kind=valt)                                 :: my_value, parent_value
267      LOGICAL                                            :: all_done
268
269      DBCSR_ASSERT(0 < first .AND. first <= heap%n)
270
271      e = first
272      all_done = .FALSE.
273      IF (e .GT. 1) THEN
274         my_value = get_value(heap, e)
275      ENDIF
276      ! Check whether we are finished, i.e,. whether the element to
277      ! bubble up is an orphan.
278      new_pos = e
279      DO WHILE (e .GT. 1 .AND. .NOT. all_done)
280         ! Switches the parent and the current element if the current
281         ! element's value is greater than the parent's value.
282         parent = get_parent(e)
283         parent_value = get_value(heap, parent)
284         IF (my_value .LT. parent_value) THEN
285            CALL dbcsr_heap_swap(heap, e, parent)
286            e = parent
287         ELSE
288            all_done = .TRUE.
289         ENDIF
290      ENDDO
291      new_pos = e
292   END SUBROUTINE bubble_up
293
294END MODULE dbcsr_min_heap
295