Skip to content

Commit fffc2c5

Browse files
committed
added: AddOrKeepEntry
1 parent 4fa0004 commit fffc2c5

File tree

1 file changed

+127
-22
lines changed

1 file changed

+127
-22
lines changed

src/ImTools/ImTools.Experimental.cs

Lines changed: 127 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -232,14 +232,14 @@ public ImMapTree<V> AddOrUpdateLeftOrRight(int key, ImMapEntry<V> entry)
232232
}
233233

234234
/// <summary>Adds the left or right branch</summary>
235-
public ImMapTree<V> AddUnsafeLeftOrRight(int key, ImMapEntry<V> entry)
235+
public ImMapTree<V> AddUnsafeLeftOrRightEntry(int key, ImMapEntry<V> entry)
236236
{
237237
if (key < Entry.Key)
238238
{
239239
var left = Left;
240240
if (left is ImMapTree<V> leftTree)
241241
{
242-
var newLeftTree = leftTree.AddUnsafeLeftOrRight(key, entry);
242+
var newLeftTree = leftTree.AddUnsafeLeftOrRightEntry(key, entry);
243243
return newLeftTree.TreeHeight == leftTree.TreeHeight
244244
? new ImMapTree<V>(Entry, newLeftTree, Right, TreeHeight)
245245
: BalanceNewLeftTree(newLeftTree);
@@ -269,7 +269,7 @@ public ImMapTree<V> AddUnsafeLeftOrRight(int key, ImMapEntry<V> entry)
269269
var right = Right;
270270
if (right is ImMapTree<V> rightTree)
271271
{
272-
var newRightTree = rightTree.AddUnsafeLeftOrRight(key, entry);
272+
var newRightTree = rightTree.AddUnsafeLeftOrRightEntry(key, entry);
273273
return newRightTree.TreeHeight == rightTree.TreeHeight
274274
? new ImMapTree<V>(Entry, Left, newRightTree, TreeHeight)
275275
: BalanceNewRightTree(newRightTree);
@@ -435,16 +435,8 @@ public ImMapTree<V> AddOrKeepLeftOrRight(int key)
435435
if (key > leftBranch.Entry.Key)
436436
{
437437
var newLeft =
438-
// 5 5
439-
// 2 ? => 3 ?
440-
// 3 2 4
441-
// 4
442-
key > leftBranch.RightEntry.Key
438+
key > leftBranch.RightEntry.Key
443439
? new ImMapTree<V>(leftBranch.RightEntry, leftBranch.Entry, new ImMapEntry<V>(key))
444-
// 5 5
445-
// 2 ? => 2.5 ?
446-
// 3 2 3
447-
// 2.5
448440
: key < leftBranch.RightEntry.Key
449441
? new ImMapTree<V>(new ImMapEntry<V>(key), leftBranch.Entry, leftBranch.RightEntry)
450442
: this;
@@ -483,16 +475,8 @@ public ImMapTree<V> AddOrKeepLeftOrRight(int key)
483475
if (key > rightBranch.Entry.Key)
484476
{
485477
var newRight =
486-
// 5 5
487-
// ? 6 => ? 8
488-
// 8 6 !
489-
// !
490-
key > rightBranch.RightEntry.Key
478+
key > rightBranch.RightEntry.Key
491479
? new ImMapTree<V>(rightBranch.RightEntry, rightBranch.Entry, new ImMapEntry<V>(key))
492-
// 5 5
493-
// ? 6 => ? 7
494-
// 8 6 8
495-
// 7
496480
: key < rightBranch.RightEntry.Key
497481
? new ImMapTree<V>(new ImMapEntry<V>(key), rightBranch.Entry, rightBranch.RightEntry)
498482
: this;
@@ -516,6 +500,99 @@ public ImMapTree<V> AddOrKeepLeftOrRight(int key)
516500
}
517501
}
518502

503+
/// <summary>Adds to the left or right branch, or keeps the un-modified map</summary>
504+
public ImMapTree<V> AddOrKeepLeftOrRightEntry(int key, ImMapEntry<V> entry)
505+
{
506+
if (key < Entry.Key)
507+
{
508+
var left = Left;
509+
if (left is ImMapTree<V> leftTree)
510+
{
511+
if (key == leftTree.Entry.Key)
512+
return this;
513+
514+
var newLeftTree = leftTree.AddOrKeepLeftOrRightEntry(key, entry);
515+
return newLeftTree == leftTree ? this
516+
: newLeftTree.TreeHeight == leftTree.TreeHeight
517+
? new ImMapTree<V>(Entry, newLeftTree, Right, TreeHeight)
518+
: BalanceNewLeftTree(newLeftTree);
519+
}
520+
521+
if (left is ImMapBranch<V> leftBranch)
522+
{
523+
if (key < leftBranch.Entry.Key)
524+
return new ImMapTree<V>(Entry,
525+
new ImMapTree<V>(leftBranch.Entry, entry, leftBranch.RightEntry),
526+
Right, TreeHeight);
527+
528+
if (key > leftBranch.Entry.Key)
529+
{
530+
var newLeft =
531+
key > leftBranch.RightEntry.Key
532+
? new ImMapTree<V>(leftBranch.RightEntry, leftBranch.Entry, entry)
533+
: key < leftBranch.RightEntry.Key
534+
? new ImMapTree<V>(entry, leftBranch.Entry, leftBranch.RightEntry)
535+
: this;
536+
537+
return new ImMapTree<V>(Entry, newLeft, Right, TreeHeight);
538+
}
539+
540+
return this;
541+
}
542+
543+
var leftLeaf = (ImMapEntry<V>)left;
544+
return key > leftLeaf.Key
545+
? new ImMapTree<V>(Entry, new ImMapBranch<V>(leftLeaf, entry), Right, 3)
546+
: key < leftLeaf.Key
547+
? new ImMapTree<V>(Entry, new ImMapBranch<V>(entry, leftLeaf), Right, 3)
548+
: this;
549+
}
550+
else
551+
{
552+
var right = Right;
553+
if (right is ImMapTree<V> rightTree)
554+
{
555+
if (key == rightTree.Entry.Key)
556+
return this;
557+
558+
// note: tree always contains left and right (for the missing leaf we have a Branch)
559+
var newRightTree = rightTree.AddOrKeepLeftOrRightEntry(key, entry);
560+
return newRightTree == rightTree ? this
561+
: newRightTree.TreeHeight == rightTree.TreeHeight
562+
? new ImMapTree<V>(Entry, Left, newRightTree, TreeHeight)
563+
: BalanceNewRightTree(newRightTree);
564+
}
565+
566+
if (right is ImMapBranch<V> rightBranch)
567+
{
568+
if (key > rightBranch.Entry.Key)
569+
{
570+
var newRight =
571+
key > rightBranch.RightEntry.Key
572+
? new ImMapTree<V>(rightBranch.RightEntry, rightBranch.Entry, entry)
573+
: key < rightBranch.RightEntry.Key
574+
? new ImMapTree<V>(entry, rightBranch.Entry, rightBranch.RightEntry)
575+
: this;
576+
577+
return new ImMapTree<V>(Entry, Left, newRight, TreeHeight);
578+
}
579+
580+
return key < rightBranch.Entry.Key
581+
? new ImMapTree<V>(Entry, Left,
582+
new ImMapTree<V>(rightBranch.Entry, entry, rightBranch.RightEntry),
583+
TreeHeight)
584+
: this;
585+
}
586+
587+
var rightLeaf = (ImMapEntry<V>)right;
588+
return key > rightLeaf.Key
589+
? new ImMapTree<V>(Entry, Left, new ImMapBranch<V>(rightLeaf, entry), 3)
590+
: key < rightLeaf.Key
591+
? new ImMapTree<V>(Entry, Left, new ImMapBranch<V>(entry, rightLeaf), 3)
592+
: this;
593+
}
594+
}
595+
519596
private ImMapTree<V> BalanceNewLeftTree(ImMapTree<V> newLeftTree)
520597
{
521598
var rightHeight = Right.Height;
@@ -727,7 +804,7 @@ public static ImMap<V> AddEntryUnsafe<V>(this ImMap<V> map, ImMapEntry<V> entry)
727804
: new ImMapTree<V>(entry, branch.Entry, branch.RightEntry)
728805
: new ImMapTree<V>(branch.Entry, entry, branch.RightEntry);
729806

730-
return ((ImMapTree<V>)map).AddUnsafeLeftOrRight(key, entry);
807+
return ((ImMapTree<V>)map).AddUnsafeLeftOrRightEntry(key, entry);
731808
}
732809

733810
/// <summary> Adds the value for the key or returns the un-modified map if key is already present </summary>
@@ -792,6 +869,34 @@ public static ImMap<V> AddOrKeep<V>(this ImMap<V> map, int key)
792869
return key != tree.Entry.Key ? tree.AddOrKeepLeftOrRight(key) : map;
793870
}
794871

872+
/// <summary> Adds the entry for the key or returns the un-modified map if key is already present </summary>
873+
[MethodImpl((MethodImplOptions)256)]
874+
public static ImMap<V> AddOrKeepEntry<V>(this ImMap<V> map, ImMapEntry<V> entry)
875+
{
876+
if (map == ImMap<V>.Empty)
877+
return entry;
878+
879+
var key = entry.Key;
880+
if (map is ImMapEntry<V> leaf)
881+
return key > leaf.Key ? new ImMapBranch<V>(leaf, entry)
882+
: key < leaf.Key ? new ImMapBranch<V>(entry, leaf)
883+
: map;
884+
885+
if (map is ImMapBranch<V> branch)
886+
{
887+
if (key > branch.Entry.Key)
888+
return key > branch.RightEntry.Key ? new ImMapTree<V>(branch.RightEntry, branch.Entry, entry)
889+
: key < branch.RightEntry.Key ? new ImMapTree<V>(entry, branch.Entry, branch.RightEntry)
890+
: map;
891+
return key < branch.Entry.Key
892+
? new ImMapTree<V>(branch.Entry, entry, branch.RightEntry)
893+
: map;
894+
}
895+
896+
var tree = (ImMapTree<V>)map;
897+
return key != tree.Entry.Key ? tree.AddOrKeepLeftOrRightEntry(key, entry) : map;
898+
}
899+
795900
///<summary>Returns the new map with the updated value for the key, or the same map if the key was not found.</summary>
796901
[MethodImpl((MethodImplOptions)256)]
797902
public static ImMap<V> Update<V>(this ImMap<V> map, int key, V value) =>

0 commit comments

Comments
 (0)