Skip to content

[CompilerPerf] Changed Map.count + Set.count from O(n) to O(1) (AVL logic based on size not height) #5365

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
manofstick wants to merge 3 commits into from
Closed

[CompilerPerf] Changed Map.count + Set.count from O(n) to O(1) (AVL logic based on size not height) #5365

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c1c30a8
Changed Map.count from O(n) to O(1)
manofstick Jul 22, 2018
63dc9b5
Modified Set to use size rather than height
manofstick Jul 23, 2018
cbc0d28
Tagged collections Map + Set count from O(n) to O(1) and rebalance vi…
manofstick Jul 24, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 20 additions & 26 deletions src/fsharp/FSharp.Core/map.fs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,22 +13,13 @@ namespace Microsoft.FSharp.Collections
type MapTree<'Key,'Value when 'Key : comparison > =
| MapEmpty
| MapOne of 'Key * 'Value
| MapNode of 'Key * 'Value * MapTree<'Key,'Value> * MapTree<'Key,'Value> * int
| MapNode of 'Key * 'Value * MapTree<'Key,'Value> * MapTree<'Key,'Value> * size:int
// REVIEW: performance rumour has it that the data held in MapNode and MapOne should be
// exactly one cache line. It is currently ~7 and 4 words respectively.

[<CompilationRepresentation(CompilationRepresentationFlags.ModuleSuffix)>]
module MapTree =

let rec sizeAux acc m =
match m with
| MapEmpty -> acc
| MapOne _ -> acc + 1
| MapNode(_,_,l,r,_) -> sizeAux (sizeAux (acc+1) l) r

let size x = sizeAux 0 x


#if TRACE_SETS_AND_MAPS
let mutable traceCount = 0
let mutable numOnes = 0
Expand Down Expand Up @@ -64,10 +55,10 @@ namespace Microsoft.FSharp.Collections

let empty = MapEmpty

let height = function
let size = function
Copy link
Contributor

@forki forki Jul 23, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please inline size function

Copy link
Contributor Author

@manofstick manofstick Jul 23, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No really keen to inline whilst MapOne exists. size as it stands has to do two attempted types casts which easily outweight a function call.

| MapEmpty -> 0
| MapOne _ -> 1
| MapNode(_,_,_,_,h) -> h
| MapNode (size=s) -> s

let isEmpty m =
match m with
Expand All @@ -78,37 +69,40 @@ namespace Microsoft.FSharp.Collections
match l,r with
| MapEmpty,MapEmpty -> MapOne(k,v)
| _ ->
let hl = height l
let hr = height r
let m = if hl < hr then hr else hl
MapNode(k,v,l,r,m+1)
let sl = size l
let sr = size r
MapNode(k,v,l,r,sl+sr+1)

let rebalance t1 k v t2 =
let t1h = height t1
let t2h = height t2
if t2h > t1h + 2 then (* right is heavier than left *)
let t1s = size t1
let t2s = size t2
if (t2s >>> 1) > t1s then (* right is over twice as heavy as left *)
match t2 with
| MapNode(t2k,t2v,t2l,t2r,_) ->
(* one of the nodes must have height > height t1 + 1 *)
if height t2l > t1h + 1 then (* balance left: combination *)
(* one of the nodes must have size > size t1 *)
if size t2l > t1s then (* balance left: combination *)
match t2l with
| MapNode(t2lk,t2lv,t2ll,t2lr,_) ->
mk (mk t1 k v t2ll) t2lk t2lv (mk t2lr t2k t2v t2r)
| _ -> failwith "rebalance"
| MapOne(t2lk,t2lv) ->
mk (mk t1 k v empty) t2lk t2lv (mk empty t2k t2v t2r)
| MapEmpty -> failwith "rebalance"
else (* rotate left *)
mk (mk t1 k v t2l) t2k t2v t2r
| _ -> failwith "rebalance"
else
if t1h > t2h + 2 then (* left is heavier than right *)
if (t1s >>> 1) > t2s then (* left is over twice as heavy as right *)
match t1 with
| MapNode(t1k,t1v,t1l,t1r,_) ->
(* one of the nodes must have height > height t2 + 1 *)
if height t1r > t2h + 1 then
(* one of the nodes must have size > size t2 *)
if size t1r > t2s then
(* balance right: combination *)
match t1r with
| MapNode(t1rk,t1rv,t1rl,t1rr,_) ->
mk (mk t1l t1k t1v t1rl) t1rk t1rv (mk t1rr k v t2)
| _ -> failwith "rebalance"
| MapOne(t1rk,t1rv) ->
mk (mk t1l t1k t1v empty) t1rk t1rv (mk empty k v t2)
| MapEmpty -> failwith "rebalance"
else
mk t1l t1k t1v (mk t1r k v t2)
| _ -> failwith "rebalance"
Expand Down
104 changes: 48 additions & 56 deletions src/fsharp/FSharp.Core/set.fs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,9 @@ namespace Microsoft.FSharp.Collections
[<CompilationRepresentation(CompilationRepresentationFlags.UseNullAsTrueValue)>]
[<NoEquality; NoComparison>]
type SetTree<'T> when 'T: comparison =
| SetEmpty // height = 0
| SetNode of 'T * SetTree<'T> * SetTree<'T> * int // height = int
| SetOne of 'T // height = 1
| SetEmpty // size = 0
| SetNode of 'T * SetTree<'T> * SetTree<'T> * size:int
| SetOne of 'T // size = 1
// OPTIMIZATION: store SetNode(k,SetEmpty,SetEmpty,1) ---> SetOne(k)
// REVIEW: performance rumour has it that the data held in SetNode and SetOne should be
// exactly one cache line on typical architectures. They are currently
Expand All @@ -27,14 +27,6 @@ namespace Microsoft.FSharp.Collections

[<CompilationRepresentation(CompilationRepresentationFlags.ModuleSuffix)>]
module internal SetTree =
let rec countAux s acc =
match s with
| SetNode(_,l,r,_) -> countAux l (countAux r (acc+1))
| SetOne(_) -> acc+1
| SetEmpty -> acc

let count s = countAux s 0

#if TRACE_SETS_AND_MAPS
let mutable traceCount = 0
let mutable numOnes = 0
Expand Down Expand Up @@ -65,16 +57,17 @@ namespace Microsoft.FSharp.Collections
n
#else
let SetOne n = SetTree.SetOne n
let SetNode (x,l,r,h) = SetTree.SetNode(x,l,r,h)
let SetNode (x,l,r,s) = SetTree.SetNode(x,l,r,s)

#endif


let height t =
let empty = SetEmpty

let size t =
match t with
| SetEmpty -> 0
| SetOne _ -> 1
| SetNode (_,_,_,h) -> h
| SetNode (size=s) -> s

#if CHECKED
let rec checkInvariant t =
Expand All @@ -88,46 +81,47 @@ namespace Microsoft.FSharp.Collections
(-2 <= (h1 - h2) && (h1 - h2) <= 2) && checkInvariant t1 && checkInvariant t2
#endif

let tolerance = 2

let mk l k r =
match l,r with
| SetEmpty,SetEmpty -> SetOne (k)
| _ ->
let hl = height l
let hr = height r
let m = if hl < hr then hr else hl
SetNode(k,l,r,m+1)
let sl = size l
let sr = size r
SetNode(k,l,r,sl+sr+1)

let rebalance t1 k t2 =
let t1h = height t1
let t2h = height t2
if t2h > t1h + tolerance then // right is heavier than left
let t1s = size t1
let t2s = size t2
if (t2s >>> 1) > t1s then (* right is over twice as heavy as left *)
match t2 with
| SetNode(t2k,t2l,t2r,_) ->
// one of the nodes must have height > height t1 + 1
if height t2l > t1h + 1 then // balance left: combination
match t2l with
| SetNode(t2lk,t2ll,t2lr,_) ->
mk (mk t1 k t2ll) t2lk (mk t2lr t2k t2r)
| _ -> failwith "rebalance"
else // rotate left
mk (mk t1 k t2l) t2k t2r
(* one of the nodes must have size > size t1 *)
if size t2l > t1s then (* balance left: combination *)
match t2l with
| SetNode(t2lk,t2ll,t2lr,_) ->
mk (mk t1 k t2ll) t2lk (mk t2lr t2k t2r)
| SetOne(t2lk) ->
mk (mk t1 k empty) t2lk (mk empty t2k t2r)
| SetEmpty -> failwith "rebalance"
else (* rotate left *)
mk (mk t1 k t2l) t2k t2r
| _ -> failwith "rebalance"
else
if t1h > t2h + tolerance then // left is heavier than right
match t1 with
| SetNode(t1k,t1l,t1r,_) ->
// one of the nodes must have height > height t2 + 1
if height t1r > t2h + 1 then
// balance right: combination
match t1r with
| SetNode(t1rk,t1rl,t1rr,_) ->
mk (mk t1l t1k t1rl) t1rk (mk t1rr k t2)
| _ -> failwith "rebalance"
else
mk t1l t1k (mk t1r k t2)
| _ -> failwith "rebalance"
if (t1s >>> 1) > t2s then (* left is over twice as heavy as right *)
match t1 with
| SetNode(t1k,t1l,t1r,_) ->
(* one of the nodes must have size > size t2 *)
if size t1r > t2s then
(* balance right: combination *)
match t1r with
| SetNode(t1rk,t1rl,t1rr,_) ->
mk (mk t1l t1k t1rl) t1rk (mk t1rr k t2)
| SetOne(t1rk) ->
mk (mk t1l t1k empty) t1rk (mk empty k t2)
| SetEmpty -> failwith "rebalance"
else
mk t1l t1k (mk t1r k t2)
| _ -> failwith "rebalance"
else mk t1 k t2

let rec add (comparer: IComparer<'T>) k t =
Expand All @@ -148,22 +142,22 @@ namespace Microsoft.FSharp.Collections
let rec balance comparer t1 k t2 =
// Given t1 < k < t2 where t1 and t2 are "balanced",
// return a balanced tree for <t1,k,t2>.
// Recall: balance means subtrees heights differ by at most "tolerance"
// Recall: balance means subtrees size of trees differ by a factor of 2
match t1,t2 with
| SetEmpty,t2 -> add comparer k t2 // drop t1 = empty
| t1,SetEmpty -> add comparer k t1 // drop t2 = empty
| SetOne k1,t2 -> add comparer k (add comparer k1 t2)
| t1,SetOne k2 -> add comparer k (add comparer k2 t1)
| SetNode(k1,t11,t12,h1),SetNode(k2,t21,t22,h2) ->
| SetNode(k1,t11,t12,s1),SetNode(k2,t21,t22,s2) ->
// Have: (t11 < k1 < t12) < k < (t21 < k2 < t22)
// Either (a) h1,h2 differ by at most 2 - no rebalance needed.
// (b) h1 too small, i.e. h1+2 < h2
// (c) h2 too small, i.e. h2+2 < h1
if h1+tolerance < h2 then
// Either (a) h1,h2 differ by at most a factory of 2 - no rebalance needed.
// (b) h1 too small
// (c) h2 too small
if (s2 >>> 1) > s1 then
// case: b, h1 too small
// push t1 into low side of t2, may increase height by 1 so rebalance
rebalance (balance comparer t1 k t21) k2 t22
elif h2+tolerance < h1 then
elif (s1 >>> 1) > s2 then
// case: c, h2 too small
// push t2 into high side of t1, may increase height by 1 so rebalance
rebalance t11 k1 (balance comparer t12 k t2)
Expand Down Expand Up @@ -491,13 +485,11 @@ namespace Microsoft.FSharp.Collections
iter (fun x -> arr.[!j] <- x; j := !j + 1) s

let toArray s =
let n = (count s)
let n = size s
let res = Array.zeroCreate n
copyToArray s res 0;
res



let rec mkFromEnumerator comparer acc (e: IEnumerator<_>) =
if e.MoveNext() then
mkFromEnumerator comparer (add comparer e.Current acc) e
Expand Down Expand Up @@ -582,7 +574,7 @@ namespace Microsoft.FSharp.Collections
#endif
Set<'T>(s.Comparer,SetTree.remove s.Comparer value s.Tree)

member s.Count = SetTree.count s.Tree
member s.Count = SetTree.size s.Tree

member s.Contains(value) =
#if TRACE_SETS_AND_MAPS
Expand Down
Loading