Skip to content

Commit 2920b00

Browse files
tiensonqintiensonqin
committed
Optimize cursor sequence traversal
1 parent cc41a19 commit 2920b00

1 file changed

Lines changed: 120 additions & 48 deletions

File tree

‎lib/persistent_sorted_set.ml‎

Lines changed: 120 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -1542,49 +1542,6 @@ let stop_cursor cursor =
15421542
cursor.stack <- [];
15431543
cursor.leaf <- None
15441544

1545-
let rec next_cursor_value_asc cursor () =
1546-
match cursor.leaf with
1547-
| None -> Seq.Nil
1548-
| Some values ->
1549-
if cursor.leaf_index >= cursor.leaf_len then (
1550-
cursor.leaf <- None;
1551-
load_next_cursor_leaf cursor;
1552-
next_cursor_value_asc cursor ())
1553-
else
1554-
let value = values.(cursor.leaf_index) in
1555-
cursor.leaf_index <- cursor.leaf_index + 1;
1556-
Seq.Cons (value, next_cursor_value_asc cursor)
1557-
1558-
let rec next_cursor_value_desc cursor () =
1559-
match cursor.leaf with
1560-
| None -> Seq.Nil
1561-
| Some values ->
1562-
if cursor.leaf_index < 0 then (
1563-
cursor.leaf <- None;
1564-
load_next_cursor_leaf cursor;
1565-
next_cursor_value_desc cursor ())
1566-
else
1567-
let value = values.(cursor.leaf_index) in
1568-
let above_upper =
1569-
match cursor.upper with
1570-
| Some upper -> cursor.cmp value upper > 0
1571-
| None -> false
1572-
in
1573-
let below_lower =
1574-
match cursor.lower with
1575-
| Some lower -> cursor.cmp value lower < 0
1576-
| None -> false
1577-
in
1578-
if above_upper then (
1579-
cursor.leaf_index <- cursor.leaf_index - 1;
1580-
next_cursor_value_desc cursor ())
1581-
else if below_lower then (
1582-
stop_cursor cursor;
1583-
Seq.Nil)
1584-
else (
1585-
cursor.leaf_index <- cursor.leaf_index - 1;
1586-
Seq.Cons (value, next_cursor_value_desc cursor))
1587-
15881545
let cursor_seq storage cmp direction lower upper root =
15891546
let cursor =
15901547
{
@@ -1601,11 +1558,112 @@ let cursor_seq storage cmp direction lower upper root =
16011558
leaf_index = 0;
16021559
}
16031560
in
1604-
fun () ->
1561+
let rec next_asc () =
16051562
init_cursor cursor;
1606-
match direction with
1607-
| Asc -> next_cursor_value_asc cursor ()
1608-
| Desc -> next_cursor_value_desc cursor ()
1563+
match cursor.leaf with
1564+
| None -> Seq.Nil
1565+
| Some values ->
1566+
if cursor.leaf_index >= cursor.leaf_len then (
1567+
cursor.leaf <- None;
1568+
load_next_cursor_leaf cursor;
1569+
next_asc ())
1570+
else
1571+
let value = values.(cursor.leaf_index) in
1572+
cursor.leaf_index <- cursor.leaf_index + 1;
1573+
Seq.Cons (value, next_asc)
1574+
in
1575+
let rec next_desc () =
1576+
init_cursor cursor;
1577+
match cursor.leaf with
1578+
| None -> Seq.Nil
1579+
| Some values ->
1580+
if cursor.leaf_index < 0 then (
1581+
cursor.leaf <- None;
1582+
load_next_cursor_leaf cursor;
1583+
next_desc ())
1584+
else
1585+
let value = values.(cursor.leaf_index) in
1586+
let above_upper =
1587+
match cursor.upper with
1588+
| Some upper -> cursor.cmp value upper > 0
1589+
| None -> false
1590+
in
1591+
let below_lower =
1592+
match cursor.lower with
1593+
| Some lower -> cursor.cmp value lower < 0
1594+
| None -> false
1595+
in
1596+
if above_upper then (
1597+
cursor.leaf_index <- cursor.leaf_index - 1;
1598+
next_desc ())
1599+
else if below_lower then (
1600+
stop_cursor cursor;
1601+
Seq.Nil)
1602+
else (
1603+
cursor.leaf_index <- cursor.leaf_index - 1;
1604+
Seq.Cons (value, next_desc))
1605+
in
1606+
match direction with Asc -> next_asc | Desc -> next_desc
1607+
1608+
let cursor_fold storage cmp direction lower upper root f init =
1609+
let cursor =
1610+
{
1611+
storage;
1612+
cmp;
1613+
direction;
1614+
lower;
1615+
upper;
1616+
root;
1617+
initialized = false;
1618+
stack = [];
1619+
leaf = None;
1620+
leaf_len = 0;
1621+
leaf_index = 0;
1622+
}
1623+
in
1624+
let rec loop acc =
1625+
init_cursor cursor;
1626+
match cursor.leaf with
1627+
| None -> acc
1628+
| Some values -> (
1629+
match cursor.direction with
1630+
| Asc ->
1631+
if cursor.leaf_index >= cursor.leaf_len then (
1632+
cursor.leaf <- None;
1633+
load_next_cursor_leaf cursor;
1634+
loop acc)
1635+
else
1636+
let value = values.(cursor.leaf_index) in
1637+
cursor.leaf_index <- cursor.leaf_index + 1;
1638+
loop (f acc value)
1639+
| Desc ->
1640+
if cursor.leaf_index < 0 then (
1641+
cursor.leaf <- None;
1642+
load_next_cursor_leaf cursor;
1643+
loop acc)
1644+
else
1645+
let value = values.(cursor.leaf_index) in
1646+
let above_upper =
1647+
match cursor.upper with
1648+
| Some upper -> cursor.cmp value upper > 0
1649+
| None -> false
1650+
in
1651+
let below_lower =
1652+
match cursor.lower with
1653+
| Some lower -> cursor.cmp value lower < 0
1654+
| None -> false
1655+
in
1656+
if above_upper then (
1657+
cursor.leaf_index <- cursor.leaf_index - 1;
1658+
loop acc)
1659+
else if below_lower then (
1660+
stop_cursor cursor;
1661+
acc)
1662+
else (
1663+
cursor.leaf_index <- cursor.leaf_index - 1;
1664+
loop (f acc value)))
1665+
in
1666+
loop init
16091667

16101668
let slice ?from_ ?to_ ?cmp (set : 'a t) =
16111669
let cmp = Option.value ~default:set.cmp cmp in
@@ -1680,7 +1738,21 @@ let seq_reverse (seq : 'a seq) =
16801738
let direction = match seq.direction with Asc -> Desc | Desc -> Asc in
16811739
{ seq with direction }
16821740

1683-
let fold_seq f init (seq : 'a seq) = Seq.fold_left f init (to_seq seq)
1741+
let fold_seq f init (seq : 'a seq) =
1742+
match (seq.source, seq.direction) with
1743+
| Seq_empty, (Asc | Desc) -> init
1744+
| Seq_tree { storage; root }, Asc ->
1745+
cursor_fold storage seq.set_cmp Asc seq.lower seq.upper (Cursor_node root)
1746+
f init
1747+
| Seq_tree { storage; root }, Desc ->
1748+
cursor_fold storage seq.set_cmp Desc seq.lower seq.upper (Cursor_node root)
1749+
f init
1750+
| Seq_deferred { storage; address }, Asc ->
1751+
cursor_fold (Some storage) seq.set_cmp Asc seq.lower seq.upper
1752+
(Cursor_address address) f init
1753+
| Seq_deferred { storage; address }, Desc ->
1754+
cursor_fold (Some storage) seq.set_cmp Desc seq.lower seq.upper
1755+
(Cursor_address address) f init
16841756

16851757
let slice_seq ?from_ ?to_ ?cmp (set : 'a t) =
16861758
let set_cmp = Option.value ~default:set.cmp cmp |> normalize_cmp in

0 commit comments

Comments
 (0)