Skip to content

Commit 0cdca5b

Browse files
AsterlessAsterless
committed
fix(HAMT): fix union_with, intersection, intersection_with and difference methods, now they can handle branchs
1 parent 1b84483 commit 0cdca5b

File tree

3 files changed

+127
-118
lines changed

3 files changed

+127
-118
lines changed

‎immut/hashmap/HAMT.mbt

Lines changed: 59 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -364,15 +364,26 @@ pub fn intersection[K : Eq + Hash, V](
364364
self : T[K, V],
365365
other : T[K, V]
366366
) -> T[K, V] {
367-
self
368-
.iter()
369-
.fold(init=Empty, fn(m, kv) {
370-
if other.get(kv.0) is Some(_) {
371-
m.add(kv.0, kv.1)
372-
} else {
373-
m
374-
}
375-
})
367+
match (self, other) {
368+
(_, Empty) => Empty
369+
(Empty, _) => Empty
370+
(Leaf(k, v), _) =>
371+
match other.get(k) {
372+
Some(_) => Leaf(k, v)
373+
None => Empty
374+
}
375+
(_, Leaf(k, _)) =>
376+
match self.get(k) {
377+
Some(v) => Leaf(k, v)
378+
None => Empty
379+
}
380+
(Branch(sa1), Branch(sa2)) =>
381+
Branch(sa1.intersection(sa2, fn(m1, m2) { m1.intersection(m2) }))
382+
(_, _) =>
383+
self.iter().fold(init=Empty, fn(m, kv) {
384+
if other.get(kv.0) is Some(_) { m.add(kv.0, kv.1) } else { m }
385+
})
386+
}
376387
}
377388

378389
///|
@@ -382,28 +393,49 @@ pub fn intersection_with[K : Eq + Hash, V](
382393
other : T[K, V],
383394
f : (K, V, V) -> V
384395
) -> T[K, V] {
385-
self
386-
.iter()
387-
.fold(init=Empty, fn(m, kv) {
388-
match other.get(kv.0) {
389-
Some(v2) => m.add(kv.0, f(kv.0, kv.1, v2))
390-
None => m
391-
}
392-
})
396+
match (self, other) {
397+
(_, Empty) => Empty
398+
(Empty, _) => Empty
399+
(Leaf(k, v), _) =>
400+
match other.get(k) {
401+
Some(v2) => Leaf(k, f(k, v, v2))
402+
None => Empty
403+
}
404+
(_, Leaf(k, v2)) =>
405+
match self.get(k) {
406+
Some(v1) => Leaf(k, f(k, v1, v2))
407+
None => Empty
408+
}
409+
(Branch(sa1), Branch(sa2)) =>
410+
Branch(sa1.intersection(sa2, fn(m1, m2) { m1.intersection_with(m2, f) }))
411+
(_, _) =>
412+
self.iter().fold(init=Empty, fn(m, kv) {
413+
match other.get(kv.0) {
414+
Some(v2) => m.add(kv.0, f(kv.0, kv.1, v2))
415+
None => m
416+
}
417+
})
418+
}
393419
}
394420

395421
///|
396-
/// Difference two hashmaps
422+
/// Difference of two hashmaps: elements in `self` but not in `other`
397423
pub fn difference[K : Eq + Hash, V](self : T[K, V], other : T[K, V]) -> T[K, V] {
398-
self
399-
.iter()
400-
.fold(init=Empty, fn(m, kv) {
401-
if other.get(kv.0) is None {
402-
m.add(kv.0, kv.1)
403-
} else {
404-
m
405-
}
406-
})
424+
match (self, other) {
425+
(Empty, _) => Empty
426+
(_, Empty) => self
427+
(Leaf(k, v), _) =>
428+
match other.get(k) {
429+
Some(_) => Empty
430+
None => Leaf(k, v)
431+
}
432+
(Branch(sa1), Branch(sa2)) =>
433+
Branch(sa1.difference(sa2))
434+
(_, _) =>
435+
self.iter().fold(init=Empty, fn(m, kv) {
436+
if other.get(kv.0) is None { m.add(kv.0, kv.1) } else { m }
437+
})
438+
}
407439
}
408440

409441
///|

‎immut/hashmap/HAMT_test.mbt

Lines changed: 59 additions & 91 deletions
Original file line numberDiff line numberDiff line change
@@ -581,128 +581,96 @@ test "HAMT::union leaf to non-overlapping map" {
581581
}
582582

583583
///|
584-
test "HAMT::union_with no overlap" {
585-
let m1 = @hashmap.of([(1, 10), (2, 20)])
586-
let m2 = @hashmap.of([(3, 30), (4, 40)])
587-
let u = m1.union_with(m2, fn(_k, v1, v2) { v1 + v2 })
588-
assert_eq!(u.get(1), Some(10))
589-
assert_eq!(u.get(2), Some(20))
590-
assert_eq!(u.get(3), Some(30))
591-
assert_eq!(u.get(4), Some(40))
592-
assert_eq!(u.size(), 4)
593-
}
594-
595-
///|
596-
test "HAMT::union_with overlap" {
597-
let m1 = @hashmap.of([(1, 10), (2, 20)])
598-
let m2 = @hashmap.of([(2, 200), (3, 300)])
599-
let u = m1.union_with(m2, fn(_k, v1, v2) { v1 * v2 })
600-
assert_eq!(u.get(1), Some(10))
601-
assert_eq!(u.get(2), Some(4000))
602-
assert_eq!(u.get(3), Some(300))
603-
assert_eq!(u.size(), 3)
604-
}
605-
606-
///|
607-
test "HAMT::union_with self" {
608-
let m = @hashmap.of([(1, 1), (2, 2), (3, 3)])
609-
let u = m.union_with(m, fn(_k, v1, v2) { v1 + v2 })
610-
assert_eq!(u.get(1), Some(2))
611-
assert_eq!(u.get(2), Some(4))
612-
assert_eq!(u.get(3), Some(6))
613-
assert_eq!(u.size(), 3)
584+
test "HAMT::intersection with empty" {
585+
let m1 = @hashmap.of([(1, 1)])
586+
let m2 = @hashmap.new()
587+
assert_eq!(m1.intersection(m2).size(), 0)
588+
assert_eq!(m2.intersection(m1).size(), 0)
614589
}
615590

616591
///|
617-
test "HAMT::intersection both empty" {
618-
let m1 = @hashmap.of([(1, 10), (2, 20), (3, 30)])
619-
let m2 = @hashmap.new()
620-
let inter = m1.intersection(m2)
621-
assert_eq!(inter.size(), 0)
592+
test "HAMT::intersection with leaf" {
593+
let m1 = @hashmap.of([(1, 1), (2, 2)])
594+
let m2 = @hashmap.singleton(2, 2)
595+
let m3 = @hashmap.singleton(3, 3)
596+
assert_eq!(m1.intersection(m2).get(2), Some(2))
597+
assert_eq!(m1.intersection(m3).get(3), None)
598+
assert_eq!(m2.intersection(m1).get(2), Some(2))
622599
}
623600

624601
///|
625-
test "HAMT::intersection partial overlap" {
626-
let m1 = @hashmap.of([(1, 10), (2, 20), (3, 30)])
627-
let m2 = @hashmap.of([(2, 200), (3, 300), (4, 400)])
602+
test "HAMT::intersection with branch" {
603+
let m1 = @hashmap.of([(1, 1), (2, 2), (3, 3)])
604+
let m2 = @hashmap.of([(2, 2), (3, 30), (4, 4)])
628605
let inter = m1.intersection(m2)
629606
assert_eq!(inter.get(1), None)
630-
assert_eq!(inter.get(2), Some(20))
631-
assert_eq!(inter.get(3), Some(30))
607+
assert_eq!(inter.get(2), Some(2))
608+
assert_eq!(inter.get(3), Some(3))
632609
assert_eq!(inter.get(4), None)
633-
assert_eq!(inter.size(), 2)
634610
}
635611

636612
///|
637-
test "HAMT::intersection full overlap" {
638-
let m1 = @hashmap.of([(1, 10), (2, 20)])
639-
let m2 = @hashmap.of([(1, 100), (2, 200)])
640-
let inter = m1.intersection(m2)
641-
assert_eq!(inter.get(1), Some(10))
642-
assert_eq!(inter.get(2), Some(20))
643-
assert_eq!(inter.size(), 2)
613+
test "HAMT::intersection_with with empty" {
614+
let m1 = @hashmap.of([(1, 1)])
615+
let m2 = @hashmap.new()
616+
assert_eq!(m1.intersection_with(m2, fn(_k, v1, v2) { v1 + v2 }).size(), 0)
617+
assert_eq!(m2.intersection_with(m1, fn(_k, v1, v2) { v1 + v2 }).size(), 0)
644618
}
645619

646620
///|
647-
test "HAMT::intersection_with no overlap" {
648-
let m1 = @hashmap.of([(1, 10)])
649-
let m2 = @hashmap.of([(2, 20)])
650-
let inter = m1.intersection_with(m2, fn(_k, v1, v2) { v1 + v2 })
651-
assert_eq!(inter.size(), 0)
621+
test "HAMT::intersection_with with leaf" {
622+
let m1 = @hashmap.of([(1, 1), (2, 2)])
623+
let m2 = @hashmap.singleton(2, 20)
624+
let m3 = @hashmap.singleton(3, 30)
625+
assert_eq!(
626+
m1.intersection_with(m2, fn(_k, v1, v2) { v1 + v2 }).get(2),
627+
Some(22),
628+
)
629+
assert_eq!(m1.intersection_with(m3, fn(_k, v1, v2) { v1 + v2 }).get(3), None)
630+
assert_eq!(
631+
m2.intersection_with(m1, fn(_k, v1, v2) { v1 + v2 }).get(2),
632+
Some(22),
633+
)
652634
}
653635

654636
///|
655-
test "HAMT::intersection_with has overlap" {
656-
let m1 = @hashmap.of([(1, 10), (2, 20)])
657-
let m2 = @hashmap.of([(2, 200), (3, 300)])
658-
let inter = m1.intersection_with(m2, fn(_k, v1, v2) { v1 + v2 })
637+
test "HAMT::intersection_with with branch" {
638+
let m1 = @hashmap.of([(1, 1), (2, 2), (3, 3)])
639+
let m2 = @hashmap.of([(2, 20), (3, 30), (4, 4)])
640+
let inter = m1.intersection_with(m2, fn(_k, v1, v2) { v1 * v2 })
659641
assert_eq!(inter.get(1), None)
660-
assert_eq!(inter.get(2), Some(220))
661-
assert_eq!(inter.get(3), None)
662-
assert_eq!(inter.size(), 1)
642+
assert_eq!(inter.get(2), Some(40))
643+
assert_eq!(inter.get(3), Some(90))
644+
assert_eq!(inter.get(4), None)
663645
}
664646

665647
///|
666-
test "HAMT::intersection_with string values" {
667-
let m1 = @hashmap.of([(1, "foo"), (2, "bar")])
668-
let m2 = @hashmap.of([(2, "baz"), (3, "qux")])
669-
let inter = m1.intersection_with(m2, fn(_k, v1, v2) { v1 + v2 })
670-
assert_eq!(inter.get(1), None)
671-
assert_eq!(inter.get(2), Some("barbaz"))
672-
assert_eq!(inter.get(3), None)
673-
assert_eq!(inter.size(), 1)
648+
test "HAMT::difference with empty" {
649+
let m1 = @hashmap.of([(1, 1)])
650+
let m2 = @hashmap.new()
651+
assert_eq!(m1.difference(m2), m1)
652+
assert_eq!(m2.difference(m1).size(), 0)
674653
}
675654

676655
///|
677-
test "HAMT::difference no overlap" {
678-
let m1 = @hashmap.of([(1, 10), (2, 20)])
679-
let m2 = @hashmap.of([(3, 30), (4, 40)])
680-
let diff = m1.difference(m2)
681-
assert_eq!(diff.get(1), Some(10))
682-
assert_eq!(diff.get(2), Some(20))
683-
assert_eq!(diff.get(3), None)
684-
assert_eq!(diff.get(4), None)
685-
assert_eq!(diff.size(), 2)
656+
test "HAMT::difference with leaf" {
657+
let m1 = @hashmap.of([(1, 1), (2, 2)])
658+
let m2 = @hashmap.singleton(2, 2)
659+
let m3 = @hashmap.singleton(3, 3)
660+
assert_eq!(m1.difference(m2).get(2), None)
661+
assert_eq!(m1.difference(m3).get(1), Some(1))
662+
assert_eq!(m2.difference(m1).size(), 0)
686663
}
687664

688665
///|
689-
test "HAMT::difference partial overlap" {
690-
let m1 = @hashmap.of([(1, 10), (2, 20), (3, 30)])
691-
let m2 = @hashmap.of([(2, 200), (4, 400)])
666+
test "HAMT::difference with branch" {
667+
let m1 = @hashmap.of([(1, 1), (2, 2), (3, 3)])
668+
let m2 = @hashmap.of([(2, 2), (3, 30), (4, 4)])
692669
let diff = m1.difference(m2)
693-
assert_eq!(diff.get(1), Some(10))
670+
assert_eq!(diff.get(1), Some(1))
694671
assert_eq!(diff.get(2), None)
695-
assert_eq!(diff.get(3), Some(30))
672+
assert_eq!(diff.get(3), None)
696673
assert_eq!(diff.get(4), None)
697-
assert_eq!(diff.size(), 2)
698-
}
699-
700-
///|
701-
test "HAMT::difference full overlap" {
702-
let m1 = @hashmap.of([(1, 10), (2, 20)])
703-
let m2 = @hashmap.of([(1, 100), (2, 200)])
704-
let diff = m1.difference(m2)
705-
assert_eq!(diff.size(), 0)
706674
}
707675

708676
///|

‎immut/hashmap/hashmap.mbti

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,12 @@ fn union[K : Eq + Hash, V](T[K, V], T[K, V]) -> T[K, V]
6363

6464
fn union_with[K : Eq + Hash, V](T[K, V], T[K, V], (K, V, V) -> V) -> T[K, V]
6565

66+
fn intersection[K : Eq + Hash, V](Self[K, V], Self[K, V]) -> Self[K, V]
67+
68+
fn intersection_with[K : Eq + Hash, V](Self[K, V], Self[K, V], (K, V, V) -> V) -> Self[K, V]
69+
70+
fn difference[K : Eq + Hash, V](Self[K, V], Self[K, V]) -> Self[K, V]
71+
6672
// Types and methods
6773
type T[K, V]
6874
impl T {
@@ -91,6 +97,9 @@ impl T {
9197
to_array[K, V](Self[K, V]) -> Array[(K, V)]
9298
union[K : Eq + Hash, V](Self[K, V], Self[K, V]) -> Self[K, V]
9399
union_with[K : Eq + Hash, V](Self[K, V], Self[K, V], (K, V, V) -> V) -> Self[K, V]
100+
intersection[K : Eq + Hash, V](Self[K, V], Self[K, V]) -> Self[K, V]
101+
intersection_with[K : Eq + Hash, V](Self[K, V], Self[K, V], (K, V, V) -> V) -> Self[K, V]
102+
difference[K : Eq + Hash, V](Self[K, V], Self[K, V]) -> Self[K, V]
94103
}
95104
impl[K : Eq + Hash, V : Eq] Eq for T[K, V]
96105
impl[K : Hash, V : Hash] Hash for T[K, V]

0 commit comments

Comments
 (0)