@@ -4,6 +4,7 @@ use std::mem;
44
55use mbe:: { SyntheticToken , SyntheticTokenId , TokenMap } ;
66use rustc_hash:: FxHashMap ;
7+ use smallvec:: SmallVec ;
78use syntax:: {
89 ast:: { self , AstNode , HasLoopBody } ,
910 match_ast, SyntaxElement , SyntaxKind , SyntaxNode , TextRange ,
@@ -292,25 +293,34 @@ pub(crate) fn reverse_fixups(
292293 token_map : & TokenMap ,
293294 undo_info : & SyntaxFixupUndoInfo ,
294295) {
295- tt. token_trees . retain ( |tt| match tt {
296- tt:: TokenTree :: Leaf ( leaf) => {
297- token_map. synthetic_token_id ( leaf. id ( ) ) . is_none ( )
298- || token_map. synthetic_token_id ( leaf. id ( ) ) != Some ( EMPTY_ID )
299- }
300- tt:: TokenTree :: Subtree ( st) => st. delimiter . map_or ( true , |d| {
301- token_map. synthetic_token_id ( d. id ) . is_none ( )
302- || token_map. synthetic_token_id ( d. id ) != Some ( EMPTY_ID )
303- } ) ,
304- } ) ;
305- tt. token_trees . iter_mut ( ) . for_each ( |tt| match tt {
306- tt:: TokenTree :: Subtree ( tt) => reverse_fixups ( tt, token_map, undo_info) ,
307- tt:: TokenTree :: Leaf ( leaf) => {
308- if let Some ( id) = token_map. synthetic_token_id ( leaf. id ( ) ) {
309- let original = & undo_info. original [ id. 0 as usize ] ;
310- * tt = tt:: TokenTree :: Subtree ( original. clone ( ) ) ;
296+ let tts = std:: mem:: take ( & mut tt. token_trees ) ;
297+ tt. token_trees = tts
298+ . into_iter ( )
299+ . filter ( |tt| match tt {
300+ tt:: TokenTree :: Leaf ( leaf) => token_map. synthetic_token_id ( leaf. id ( ) ) != Some ( EMPTY_ID ) ,
301+ tt:: TokenTree :: Subtree ( st) => {
302+ st. delimiter . map_or ( true , |d| token_map. synthetic_token_id ( d. id ) != Some ( EMPTY_ID ) )
311303 }
312- }
313- } ) ;
304+ } )
305+ . flat_map ( |tt| match tt {
306+ tt:: TokenTree :: Subtree ( mut tt) => {
307+ reverse_fixups ( & mut tt, token_map, undo_info) ;
308+ SmallVec :: from_const ( [ tt. into ( ) ] )
309+ }
310+ tt:: TokenTree :: Leaf ( leaf) => {
311+ if let Some ( id) = token_map. synthetic_token_id ( leaf. id ( ) ) {
312+ let original = undo_info. original [ id. 0 as usize ] . clone ( ) ;
313+ if original. delimiter . is_none ( ) {
314+ original. token_trees . into ( )
315+ } else {
316+ SmallVec :: from_const ( [ original. into ( ) ] )
317+ }
318+ } else {
319+ SmallVec :: from_const ( [ leaf. into ( ) ] )
320+ }
321+ }
322+ } )
323+ . collect ( ) ;
314324}
315325
316326#[ cfg( test) ]
@@ -319,6 +329,31 @@ mod tests {
319329
320330 use super :: reverse_fixups;
321331
332+ // The following three functions are only meant to check partial structural equivalence of
333+ // `TokenTree`s, see the last assertion in `check()`.
334+ fn check_leaf_eq ( a : & tt:: Leaf , b : & tt:: Leaf ) -> bool {
335+ match ( a, b) {
336+ ( tt:: Leaf :: Literal ( a) , tt:: Leaf :: Literal ( b) ) => a. text == b. text ,
337+ ( tt:: Leaf :: Punct ( a) , tt:: Leaf :: Punct ( b) ) => a. char == b. char ,
338+ ( tt:: Leaf :: Ident ( a) , tt:: Leaf :: Ident ( b) ) => a. text == b. text ,
339+ _ => false ,
340+ }
341+ }
342+
343+ fn check_subtree_eq ( a : & tt:: Subtree , b : & tt:: Subtree ) -> bool {
344+ a. delimiter . map ( |it| it. kind ) == b. delimiter . map ( |it| it. kind )
345+ && a. token_trees . len ( ) == b. token_trees . len ( )
346+ && a. token_trees . iter ( ) . zip ( & b. token_trees ) . all ( |( a, b) | check_tt_eq ( a, b) )
347+ }
348+
349+ fn check_tt_eq ( a : & tt:: TokenTree , b : & tt:: TokenTree ) -> bool {
350+ match ( a, b) {
351+ ( tt:: TokenTree :: Leaf ( a) , tt:: TokenTree :: Leaf ( b) ) => check_leaf_eq ( a, b) ,
352+ ( tt:: TokenTree :: Subtree ( a) , tt:: TokenTree :: Subtree ( b) ) => check_subtree_eq ( a, b) ,
353+ _ => false ,
354+ }
355+ }
356+
322357 #[ track_caller]
323358 fn check ( ra_fixture : & str , mut expect : Expect ) {
324359 let parsed = syntax:: SourceFile :: parse ( ra_fixture) ;
@@ -331,27 +366,28 @@ mod tests {
331366 fixups. append ,
332367 ) ;
333368
334- let mut actual = tt. to_string ( ) ;
335- actual. push ( '\n' ) ;
369+ let actual = format ! ( "{}\n " , tt) ;
336370
337371 expect. indent ( false ) ;
338372 expect. assert_eq ( & actual) ;
339373
340374 // the fixed-up tree should be syntactically valid
341375 let ( parse, _) = mbe:: token_tree_to_syntax_node ( & tt, :: mbe:: TopEntryPoint :: MacroItems ) ;
342- assert_eq ! (
343- parse. errors( ) ,
344- & [ ] ,
376+ assert ! (
377+ parse. errors( ) . is_empty( ) ,
345378 "parse has syntax errors. parse tree:\n {:#?}" ,
346379 parse. syntax_node( )
347380 ) ;
348381
349382 reverse_fixups ( & mut tt, & tmap, & fixups. undo_info ) ;
350383
351384 // the fixed-up + reversed version should be equivalent to the original input
352- // (but token IDs don't matter)
385+ // modulo token IDs and `Punct`s' spacing.
353386 let ( original_as_tt, _) = mbe:: syntax_node_to_token_tree ( & parsed. syntax_node ( ) ) ;
354- assert_eq ! ( tt. to_string( ) , original_as_tt. to_string( ) ) ;
387+ assert ! (
388+ check_subtree_eq( & tt, & original_as_tt) ,
389+ "different token tree: {tt:?}, {original_as_tt:?}"
390+ ) ;
355391 }
356392
357393 #[ test]
@@ -468,7 +504,7 @@ fn foo() {
468504}
469505"# ,
470506 expect ! [ [ r#"
471- fn foo () {a .__ra_fixup}
507+ fn foo () {a . __ra_fixup}
472508"# ] ] ,
473509 )
474510 }
@@ -482,7 +518,7 @@ fn foo() {
482518}
483519"# ,
484520 expect ! [ [ r#"
485- fn foo () {a .__ra_fixup ;}
521+ fn foo () {a . __ra_fixup ;}
486522"# ] ] ,
487523 )
488524 }
@@ -497,7 +533,7 @@ fn foo() {
497533}
498534"# ,
499535 expect ! [ [ r#"
500- fn foo () {a .__ra_fixup ; bar () ;}
536+ fn foo () {a . __ra_fixup ; bar () ;}
501537"# ] ] ,
502538 )
503539 }
@@ -525,7 +561,7 @@ fn foo() {
525561}
526562"# ,
527563 expect ! [ [ r#"
528- fn foo () {let x = a .__ra_fixup ;}
564+ fn foo () {let x = a . __ra_fixup ;}
529565"# ] ] ,
530566 )
531567 }
@@ -541,7 +577,7 @@ fn foo() {
541577}
542578"# ,
543579 expect ! [ [ r#"
544- fn foo () {a .b ; bar () ;}
580+ fn foo () {a . b ; bar () ;}
545581"# ] ] ,
546582 )
547583 }
0 commit comments