Add an option for excluding instances from moves
[ganeti-local] / Ganeti / HTools / PeerMap.hs
1 {-|
2   Module abstracting the peer map implementation.
3
4 This is abstracted separately since the speed of peermap updates can
5 be a significant part of the total runtime, and as such changing the
6 implementation should be easy in case it's needed.
7
8 -}
9
10 {-
11
12 Copyright (C) 2009 Google Inc.
13
14 This program is free software; you can redistribute it and/or modify
15 it under the terms of the GNU General Public License as published by
16 the Free Software Foundation; either version 2 of the License, or
17 (at your option) any later version.
18
19 This program is distributed in the hope that it will be useful, but
20 WITHOUT ANY WARRANTY; without even the implied warranty of
21 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
22 General Public License for more details.
23
24 You should have received a copy of the GNU General Public License
25 along with this program; if not, write to the Free Software
26 Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
27 02110-1301, USA.
28
29 -}
30
31 module Ganeti.HTools.PeerMap
32     ( PeerMap
33     , Key
34     , Elem
35     , empty
36     , accumArray
37     , Ganeti.HTools.PeerMap.find
38     , add
39     , remove
40     , maxElem
41     ) where
42
43 import Data.Maybe (fromMaybe)
44 import Data.List
45 import Data.Function
46 import Data.Ord
47
48 import Ganeti.HTools.Types
49
50 type Key = Ndx
51 type Elem = Int
52 type PeerMap = [(Key, Elem)]
53
54 -- * Initialization functions
55
56 -- | Create a new empty map.
57 empty :: PeerMap
58 empty = []
59
60 -- | Our reverse-compare function.
61 pmCompare :: (Key, Elem) -> (Key, Elem) -> Ordering
62 pmCompare a b = (compare `on` snd) b a
63
64 -- | Add or update (via a custom function) an element.
65 addWith :: (Elem -> Elem -> Elem) -> Key -> Elem -> PeerMap -> PeerMap
66 addWith fn k v lst =
67     case lookup k lst of
68       Nothing -> insertBy pmCompare (k, v) lst
69       Just o -> insertBy pmCompare (k, fn o v) (remove k lst)
70
71 -- | Create a PeerMap from an association list, with possible duplicates
72 accumArray :: (Elem -> Elem -> Elem) -- ^ function used to merge the elements
73               -> [(Key, Elem)]       -- ^ source data
74               -> PeerMap             -- ^ results
75 accumArray _  [] = empty
76 accumArray fn ((k, v):xs) = addWith fn k v $ accumArray fn xs
77
78 -- * Basic operations
79
80 -- | Returns either the value for a key or zero if not found
81 find :: Key -> PeerMap -> Elem
82 find k = fromMaybe 0 . lookup k
83
84 -- | Add an element to a peermap, overwriting the previous value
85 add :: Key -> Elem -> PeerMap -> PeerMap
86 add = addWith (flip const)
87
88 -- | Remove an element from a peermap
89 remove :: Key -> PeerMap -> PeerMap
90 remove _ [] = []
91 remove k ((x@(x', _)):xs) = if k == x'
92                             then xs
93                             else x:remove k xs
94
95 -- | Find the maximum element.
96 --
97 -- Since this is a sorted list, we just get the value at the head of
98 -- the list, or zero for a null list
99 maxElem :: PeerMap -> Elem
100 maxElem (x:_) = snd x
101 maxElem _ = 0