Skip to content

Commit 6a4dd2c

Browse files
blagininblaginin
committed
Rewrite all other methods in DynTreeNode to be iterative
1 parent c26e07f commit 6a4dd2c

File tree

1 file changed

+123
-93
lines changed

1 file changed

+123
-93
lines changed

‎datafusion/common/src/tree_node.rs‎

Lines changed: 123 additions & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -912,6 +912,104 @@ macro_rules! map_until_stop_and_collect {
912912
}}
913913
}
914914

915+
macro_rules! rewrite_recursive {
916+
($START:ident, $NAME:ident, $TRANSFORM_UP:expr, $TRANSFORM_DOWN:expr) => {
917+
let mut queue = vec![ProcessingState::NotStarted($START)];
918+
919+
while let Some(item) = queue.pop() {
920+
match item {
921+
ProcessingState::NotStarted($NAME) => {
922+
let node = $TRANSFORM_DOWN?;
923+
924+
queue.push(match node.tnr {
925+
TreeNodeRecursion::Continue => {
926+
ProcessingState::ProcessingChildren {
927+
non_processed_children: node
928+
.data
929+
.arc_children()
930+
.into_iter()
931+
.cloned()
932+
.rev()
933+
.collect(),
934+
item: node,
935+
processed_children: vec![],
936+
}
937+
}
938+
TreeNodeRecursion::Jump => ProcessingState::ProcessedAllChildren(
939+
node.with_tnr(TreeNodeRecursion::Continue),
940+
),
941+
TreeNodeRecursion::Stop => {
942+
ProcessingState::ProcessedAllChildren(node)
943+
}
944+
})
945+
}
946+
ProcessingState::ProcessingChildren {
947+
mut item,
948+
mut non_processed_children,
949+
mut processed_children,
950+
} => match item.tnr {
951+
TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {
952+
if let Some(non_processed_item) = non_processed_children.pop() {
953+
queue.push(ProcessingState::ProcessingChildren {
954+
item,
955+
non_processed_children,
956+
processed_children,
957+
});
958+
queue.push(ProcessingState::NotStarted(non_processed_item));
959+
} else {
960+
item.transformed |=
961+
processed_children.iter().any(|item| item.transformed);
962+
item.data = item.data.with_new_arc_children(
963+
processed_children.into_iter().map(|c| c.data).collect(),
964+
)?;
965+
queue.push(ProcessingState::ProcessedAllChildren(item))
966+
}
967+
}
968+
TreeNodeRecursion::Stop => {
969+
processed_children.extend(
970+
non_processed_children
971+
.into_iter()
972+
.rev()
973+
.map(Transformed::no),
974+
);
975+
item.transformed |=
976+
processed_children.iter().any(|item| item.transformed);
977+
item.data = item.data.with_new_arc_children(
978+
processed_children.into_iter().map(|c| c.data).collect(),
979+
)?;
980+
queue.push(ProcessingState::ProcessedAllChildren(item));
981+
}
982+
},
983+
ProcessingState::ProcessedAllChildren(node) => {
984+
let node = node.transform_parent(|$NAME| $TRANSFORM_UP)?;
985+
986+
if let Some(ProcessingState::ProcessingChildren {
987+
item: mut parent_node,
988+
non_processed_children,
989+
mut processed_children,
990+
..
991+
}) = queue.pop()
992+
{
993+
parent_node.tnr = node.tnr;
994+
processed_children.push(node);
995+
996+
queue.push(ProcessingState::ProcessingChildren {
997+
item: parent_node,
998+
non_processed_children,
999+
processed_children,
1000+
})
1001+
} else {
1002+
debug_assert_eq!(queue.len(), 0);
1003+
return Ok(node);
1004+
}
1005+
}
1006+
}
1007+
}
1008+
1009+
unreachable!();
1010+
};
1011+
}
1012+
9151013
/// Transformation helper to access [`Transformed`] fields in a [`Result`] easily.
9161014
///
9171015
/// # Example
@@ -999,103 +1097,35 @@ impl<T: DynTreeNode + ?Sized> TreeNode for Arc<T> {
9991097
}
10001098
}
10011099

1002-
fn rewrite<R: TreeNodeRewriter<Node = Self>>(
1100+
fn transform_down_up<
1101+
FD: FnMut(Self) -> Result<Transformed<Self>>,
1102+
FU: FnMut(Self) -> Result<Transformed<Self>>,
1103+
>(
10031104
self,
1004-
rewriter: &mut R,
1105+
mut f_down: FD,
1106+
mut f_up: FU,
10051107
) -> Result<Transformed<Self>> {
1006-
let mut queue = vec![ProcessingState::NotStarted(self)];
1007-
1008-
while let Some(item) = queue.pop() {
1009-
match item {
1010-
ProcessingState::NotStarted(node) => {
1011-
let node = rewriter.f_down(node)?;
1012-
1013-
queue.push(match node.tnr {
1014-
TreeNodeRecursion::Continue => {
1015-
ProcessingState::ProcessingChildren {
1016-
non_processed_children: node
1017-
.data
1018-
.arc_children()
1019-
.into_iter()
1020-
.cloned()
1021-
.rev()
1022-
.collect(),
1023-
item: node,
1024-
processed_children: vec![],
1025-
}
1026-
}
1027-
TreeNodeRecursion::Jump => ProcessingState::ProcessedAllChildren(
1028-
node.with_tnr(TreeNodeRecursion::Continue),
1029-
),
1030-
TreeNodeRecursion::Stop => {
1031-
ProcessingState::ProcessedAllChildren(node)
1032-
}
1033-
})
1034-
}
1035-
ProcessingState::ProcessingChildren {
1036-
mut item,
1037-
mut non_processed_children,
1038-
mut processed_children,
1039-
} => match item.tnr {
1040-
TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {
1041-
if let Some(non_processed_item) = non_processed_children.pop() {
1042-
queue.push(ProcessingState::ProcessingChildren {
1043-
item,
1044-
non_processed_children,
1045-
processed_children,
1046-
});
1047-
queue.push(ProcessingState::NotStarted(non_processed_item));
1048-
} else {
1049-
item.transformed =
1050-
processed_children.iter().any(|item| item.transformed);
1051-
item.data = item.data.with_new_arc_children(
1052-
processed_children.into_iter().map(|c| c.data).collect(),
1053-
)?;
1054-
queue.push(ProcessingState::ProcessedAllChildren(item))
1055-
}
1056-
}
1057-
TreeNodeRecursion::Stop => {
1058-
processed_children.extend(
1059-
non_processed_children
1060-
.into_iter()
1061-
.rev()
1062-
.map(Transformed::no),
1063-
);
1064-
item.transformed =
1065-
processed_children.iter().any(|item| item.transformed);
1066-
item.data = item.data.with_new_arc_children(
1067-
processed_children.into_iter().map(|c| c.data).collect(),
1068-
)?;
1069-
queue.push(ProcessingState::ProcessedAllChildren(item));
1070-
}
1071-
},
1072-
ProcessingState::ProcessedAllChildren(node) => {
1073-
let node = node.transform_parent(|n| rewriter.f_up(n))?;
1074-
1075-
if let Some(ProcessingState::ProcessingChildren {
1076-
item: mut parent_node,
1077-
non_processed_children,
1078-
mut processed_children,
1079-
..
1080-
}) = queue.pop()
1081-
{
1082-
parent_node.tnr = node.tnr;
1083-
processed_children.push(node);
1108+
rewrite_recursive!(self, node, f_up(node), f_down(node));
1109+
}
10841110

1085-
queue.push(ProcessingState::ProcessingChildren {
1086-
item: parent_node,
1087-
non_processed_children,
1088-
processed_children,
1089-
})
1090-
} else {
1091-
debug_assert_eq!(queue.len(), 0);
1092-
return Ok(node);
1093-
}
1094-
}
1095-
}
1096-
}
1111+
fn transform_down<F: FnMut(Self) -> Result<Transformed<Self>>>(
1112+
self,
1113+
f: F,
1114+
) -> Result<Transformed<Self>> {
1115+
self.transform_down_up(f, |node| Ok(Transformed::no(node)))
1116+
}
10971117

1098-
unreachable!();
1118+
fn transform_up<F: FnMut(Self) -> Result<Transformed<Self>>>(
1119+
self,
1120+
f: F,
1121+
) -> Result<Transformed<Self>> {
1122+
self.transform_down_up(|node| Ok(Transformed::no(node)), f)
1123+
}
1124+
fn rewrite<R: TreeNodeRewriter<Node = Self>>(
1125+
self,
1126+
rewriter: &mut R,
1127+
) -> Result<Transformed<Self>> {
1128+
rewrite_recursive!(self, node, rewriter.f_up(node), rewriter.f_down(node));
10991129
}
11001130

11011131
fn visit<'n, V: TreeNodeVisitor<'n, Node = Self>>(

0 commit comments

Comments
 (0)