diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 39d009d..9225c6c 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -20,12 +20,18 @@ jobs:
       - name: build
         run: cargo build
 
-      - name: test
-        run: cargo test -p rstml --features "rawtext-stable-hack-module"
-
       - name: clippy
         run: cargo clippy --workspace
 
+      - name: test on Stable
+        run: cargo test --workspace
+
+      - name: Tests with rawtext hack
+        run: cargo test -p rstml --features "rawtext-stable-hack-module"
+
+      - name: Test extendable feature in rstml-control-flow
+        run: cargo test -p rstml-control-flow --features "extendable"
+
       - uses: dtolnay/rust-toolchain@nightly
 
       - name: test on Nightly
diff --git a/Cargo.toml b/Cargo.toml
index 342dcf5..0dd25c3 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -39,7 +39,7 @@ harness = false
 path = "benches/bench.rs"
 
 [workspace]
-members = ["examples/html-to-string-macro", "rstml-controll-flow"]
+members = ["examples/html-to-string-macro", "rstml-control-flow"]
 
 [features]
 default = ["colors"]
diff --git a/examples/html-to-string-macro/Cargo.toml b/examples/html-to-string-macro/Cargo.toml
index 8de1a93..7904171 100644
--- a/examples/html-to-string-macro/Cargo.toml
+++ b/examples/html-to-string-macro/Cargo.toml
@@ -22,7 +22,7 @@ rstml = { version = "0.11.0", path = "../../", features = [
 ] }
 proc-macro2-diagnostics = "0.10"
 derive-where = "1.2.5"
-rstml-controll-flow = { version = "0.1.0", path = "../../rstml-controll-flow" }
+rstml-control-flow = { version = "0.1.0", path = "../../rstml-control-flow" }
 
 [dev-dependencies]
 trybuild = "1.0"
diff --git a/examples/html-to-string-macro/src/lib.rs b/examples/html-to-string-macro/src/lib.rs
index 6124c4e..c805cdf 100644
--- a/examples/html-to-string-macro/src/lib.rs
+++ b/examples/html-to-string-macro/src/lib.rs
@@ -5,7 +5,7 @@ use quote::{quote, quote_spanned, ToTokens};
 use rstml::{
     node::{Node, NodeAttribute, NodeName},
     visitor::{visit_attributes, visit_nodes, Visitor},
-    Infallible, Parser, ParserConfig,
+    Parser, ParserConfig,
 };
 use syn::spanned::Spanned;
 // mod escape;
@@ -46,9 +46,10 @@ impl WalkNodesOutput {
 }
 impl<'a> syn::visit_mut::VisitMut for WalkNodes<'a> {}
 
-impl<'a> Visitor for WalkNodes<'a> {
-    type Custom = Infallible;
-
+impl<'a, C> Visitor<C> for WalkNodes<'a>
+where
+    C: rstml::node::CustomNode + 'static,
+{
     fn visit_doctype(&mut self, doctype: &mut rstml::node::NodeDoctype) -> bool {
         let value = &doctype.value.to_token_stream_string();
         self.output
@@ -67,7 +68,7 @@ impl<'a> Visitor for WalkNodes<'a> {
         self.output.static_format.push_str(&node.to_string_best());
         false
     }
-    fn visit_fragment(&mut self, fragment: &mut rstml::node::NodeFragment<Self::Custom>) -> bool {
+    fn visit_fragment(&mut self, fragment: &mut rstml::node::NodeFragment<C>) -> bool {
         let visitor = self.child_output();
         let child_output = visit_nodes(&mut fragment.children, visitor);
         self.output.extend(child_output.output);
@@ -85,7 +86,7 @@ impl<'a> Visitor for WalkNodes<'a> {
         self.output.values.push(block.to_token_stream());
         false
     }
-    fn visit_element(&mut self, element: &mut rstml::node::NodeElement<Self::Custom>) -> bool {
+    fn visit_element(&mut self, element: &mut rstml::node::NodeElement<C>) -> bool {
         let name = element.name().to_string();
         self.output.static_format.push_str(&format!("<{}", name));
         self.output
diff --git a/rstml-controll-flow/Cargo.toml b/rstml-control-flow/Cargo.toml
similarity index 94%
rename from rstml-controll-flow/Cargo.toml
rename to rstml-control-flow/Cargo.toml
index 08df904..f1d7bad 100644
--- a/rstml-controll-flow/Cargo.toml
+++ b/rstml-control-flow/Cargo.toml
@@ -1,5 +1,5 @@
 [package]
-name = "rstml-controll-flow"
+name = "rstml-control-flow"
 version = "0.1.0"
 edition = "2021"
 
diff --git a/rstml-controll-flow/README.md b/rstml-control-flow/README.md
similarity index 94%
rename from rstml-controll-flow/README.md
rename to rstml-control-flow/README.md
index b0f1959..0458a3b 100644
--- a/rstml-controll-flow/README.md
+++ b/rstml-control-flow/README.md
@@ -83,7 +83,7 @@ let nodes = parse2_with_config(template, Default::default().with_custom_nodes::<
 ## Using multiple `CustomNode`s at once
 It is also possible to use more than one `CustomNode` at once.
 For example, if you want to use both `Conditions` and `EscapedCode` custom nodes.
-`rstml-controll-flow` crate provides an `ExtendableCustomNode` struct that can be used to combine multiple `CustomNode`s into one. Check out `extendable.rs` docs and tests in `lib.rs` for more details.
+`rstml-control-flow` crate provides an `ExtendableCustomNode` struct that can be used to combine multiple `CustomNode`s into one. Check out `extendable.rs` docs and tests in `lib.rs` for more details.
 
 
 ```rust
\ No newline at end of file
diff --git a/rstml-controll-flow/src/escape.rs b/rstml-control-flow/src/escape.rs
similarity index 78%
rename from rstml-controll-flow/src/escape.rs
rename to rstml-control-flow/src/escape.rs
index 537b6e7..b3da803 100644
--- a/rstml-controll-flow/src/escape.rs
+++ b/rstml-control-flow/src/escape.rs
@@ -18,7 +18,6 @@ use syn::{
     braced,
     parse::{Parse, ParseStream},
     token::Brace,
-    visit_mut::VisitMut,
     Expr, Pat, Token,
 };
 
@@ -311,29 +310,47 @@ where
 }
 
 pub mod visitor_impl {
+    use rstml::visitor::{CustomNodeWalker, RustCode};
+
     use super::*;
 
+    pub struct EscapeCodeWalker;
+    impl CustomNodeWalker for EscapeCodeWalker {
+        type Custom = CustomNodeType;
+        fn walk_custom_node_fields<VisitorImpl>(
+            visitor: &mut VisitorImpl,
+            node: &mut Self::Custom,
+        ) -> bool
+        where
+            VisitorImpl: Visitor<CustomNodeType>,
+        {
+            println!("walk_custom_node_fields");
+            EscapeCode::visit_custom_children(visitor, node)
+        }
+    }
+
     impl Block {
-        pub fn visit_custom_children<V: Visitor<Custom = CustomNodeType> + VisitMut>(
+        pub fn visit_custom_children<V: Visitor<CustomNodeType>>(
             visitor: &mut V,
             block: &mut Self,
         ) -> bool {
+            println!("visit custom block {:?}", &block.body);
             block.body.iter_mut().all(|val| visitor.visit_node(val))
         }
     }
 
     impl ElseIf {
-        pub fn visit_custom_children<V: Visitor<Custom = CustomNodeType> + VisitMut>(
+        pub fn visit_custom_children<V: Visitor<CustomNodeType>>(
             visitor: &mut V,
             expr: &mut Self,
         ) -> bool {
-            visitor.visit_expr_mut(&mut expr.condition);
+            visitor.visit_rust_code(RustCode::Expr(&mut expr.condition));
             Block::visit_custom_children(visitor, &mut expr.then_branch)
         }
     }
 
     impl Else {
-        pub fn visit_custom_children<V: Visitor<Custom = CustomNodeType> + VisitMut>(
+        pub fn visit_custom_children<V: Visitor<CustomNodeType>>(
             visitor: &mut V,
             expr: &mut Self,
         ) -> bool {
@@ -342,11 +359,11 @@ pub mod visitor_impl {
     }
 
     impl IfExpr {
-        pub fn visit_custom_children<V: Visitor<Custom = CustomNodeType> + VisitMut>(
+        pub fn visit_custom_children<V: Visitor<CustomNodeType>>(
             visitor: &mut V,
             expr: &mut Self,
         ) -> bool {
-            visitor.visit_expr_mut(&mut expr.condition);
+            visitor.visit_rust_code(RustCode::Expr(&mut expr.condition));
             Block::visit_custom_children(visitor, &mut expr.then_branch)
                 || expr
                     .else_ifs
@@ -360,30 +377,30 @@ pub mod visitor_impl {
         }
     }
     impl ForExpr {
-        pub fn visit_custom_children<V: Visitor<Custom = CustomNodeType> + VisitMut>(
+        pub fn visit_custom_children<V: Visitor<CustomNodeType>>(
             visitor: &mut V,
             expr: &mut Self,
         ) -> bool {
-            visitor.visit_pat_mut(&mut expr.pat);
-            visitor.visit_expr_mut(&mut expr.expr);
+            visitor.visit_rust_code(RustCode::Pat(&mut expr.pat));
+            visitor.visit_rust_code(RustCode::Expr(&mut expr.expr));
             Block::visit_custom_children(visitor, &mut expr.block)
         }
     }
     impl Arm {
-        pub fn visit_custom_children<V: Visitor<Custom = CustomNodeType> + VisitMut>(
+        pub fn visit_custom_children<V: Visitor<CustomNodeType>>(
             visitor: &mut V,
             expr: &mut Self,
         ) -> bool {
-            visitor.visit_pat_mut(&mut expr.pat);
+            visitor.visit_rust_code(RustCode::Pat(&mut expr.pat));
             Block::visit_custom_children(visitor, &mut expr.body)
         }
     }
     impl MatchExpr {
-        pub fn visit_custom_children<V: Visitor<Custom = CustomNodeType> + VisitMut>(
+        pub fn visit_custom_children<V: Visitor<CustomNodeType>>(
             visitor: &mut V,
             expr: &mut Self,
         ) -> bool {
-            visitor.visit_expr_mut(&mut expr.expr);
+            visitor.visit_rust_code(RustCode::Expr(&mut expr.expr));
 
             expr.arms
                 .iter_mut()
@@ -391,7 +408,7 @@ pub mod visitor_impl {
         }
     }
     impl EscapedExpr {
-        pub fn visit_custom_children<V: Visitor<Custom = CustomNodeType> + VisitMut>(
+        pub fn visit_custom_children<V: Visitor<CustomNodeType>>(
             visitor: &mut V,
             expr: &mut Self,
         ) -> bool {
@@ -403,12 +420,14 @@ pub mod visitor_impl {
         }
     }
     impl EscapeCode {
-        pub fn visit_custom_children<V: Visitor<Custom = CustomNodeType> + VisitMut>(
+        pub fn visit_custom_children<V: Visitor<CustomNodeType>>(
             visitor: &mut V,
             node: &mut CustomNodeType,
         ) -> bool {
+            dbg!("visit_custom_children");
             let Either::A(mut this): Either<EscapeCode, _> = node.clone().try_into_or_clone_ref()
             else {
+                dbg!("Not a");
                 return true;
             };
             let result = EscapedExpr::visit_custom_children(visitor, &mut this.expression);
@@ -420,7 +439,7 @@ pub mod visitor_impl {
 
 #[cfg(test)]
 #[cfg(not(feature = "extendable"))]
-mod test {
+mod test_typed {
     use quote::quote;
     use rstml::{node::Node, recoverable::Recoverable, Parser, ParserConfig};
     use syn::{parse_quote, Token};
@@ -602,32 +621,6 @@ mod test {
         assert_eq!(expr.arms[1].pat, parse_quote!(y | z));
     }
 
-    #[test]
-    fn custom_node_using_config() {
-        let actual = Parser::new(
-            ParserConfig::new()
-                .element_close_use_default_wildcard_ident(false)
-                .custom_node::<MyCustomNode>(),
-        )
-        .parse_simple(quote! {
-            @if just && an || expression {
-                <a regular="html" component/>
-                <div>
-                </div>
-            }
-        })
-        .unwrap();
-        let Node::Custom(actual) = &actual[0] else {
-            panic!()
-        };
-
-        let EscapedExpr::If(expr) = &actual.expression else {
-            panic!()
-        };
-
-        assert_eq!(expr.condition, parse_quote!(just && an || expression));
-    }
-
     #[test]
     fn check_if_inside_if() {
         let actual: Recoverable<MyNode> = parse_quote! {
@@ -680,4 +673,146 @@ mod test {
         assert_eq!(expr.pat, parse_quote!(x));
         assert_eq!(expr.expr, parse_quote!(foo));
     }
+
+    #[test]
+    fn custom_node_using_config() {
+        let actual = Parser::new(
+            ParserConfig::new()
+                .element_close_use_default_wildcard_ident(false)
+                .custom_node::<MyCustomNode>(),
+        )
+        .parse_simple(quote! {
+            @if just && an || expression {
+                <a regular="html" component/>
+                <div>
+                </div>
+            }
+        })
+        .unwrap();
+        let Node::Custom(actual) = &actual[0] else {
+            panic!()
+        };
+
+        let EscapedExpr::If(expr) = &actual.expression else {
+            panic!()
+        };
+
+        assert_eq!(expr.condition, parse_quote!(just && an || expression));
+    }
+}
+
+#[cfg(test)]
+mod test_universal {
+    use proc_macro2::TokenStream;
+    use quote::quote;
+    use rstml::{
+        recoverable::Recoverable, visitor::visit_nodes_with_custom, ParserConfig, ParsingResult,
+    };
+    use syn::{parse_quote, visit_mut::VisitMut};
+
+    use super::*;
+    use crate::escape::visitor_impl::EscapeCodeWalker;
+
+    // #[cfg(feature="extendable")]
+    #[cfg(not(feature = "extendable"))]
+    fn parse_universal(input: TokenStream) -> ParsingResult<Vec<Node>> {
+        use rstml::Parser;
+
+        let actual =
+            Parser::new(ParserConfig::new().custom_node::<EscapeCode>()).parse_recoverable(input);
+
+        return actual;
+    }
+    #[cfg(feature = "extendable")]
+    fn parse_universal(input: TokenStream) -> ParsingResult<Vec<Node>> {
+        use crate::ExtendableCustomNode;
+
+        let result =
+            ExtendableCustomNode::parse2_with_config::<(EscapeCode,)>(ParserConfig::new(), input);
+        return result;
+    }
+
+    // For extendable custom node it is safe to use only after parsing.
+    fn reparse_concrete<A: ToTokens>(val: A) -> EscapeCode {
+        let recoverable: Recoverable<_> = parse_quote!(#val);
+        recoverable.inner()
+    }
+
+    #[test]
+    fn if_node_reparsable() {
+        let tokens = quote! {
+            @if just && an || expression {
+                <div/>
+            }
+        };
+
+        let actual = &parse_universal(tokens).into_result().unwrap()[0];
+
+        let Node::Custom(actual) = actual else {
+            panic!()
+        };
+        let actual = reparse_concrete(actual);
+
+        let EscapedExpr::If(expr) = actual.expression else {
+            panic!()
+        };
+
+        assert_eq!(expr.condition, parse_quote!(just && an || expression));
+    }
+    #[test]
+    fn if_node_visitor_rstml() {
+        let tokens = quote! {
+            @if just && an || expression {
+                <div/>
+            }
+        };
+
+        let mut actual = parse_universal(tokens).into_result().unwrap();
+
+        struct ElementVisitor {
+            elements: Vec<String>,
+        }
+        impl<C: std::fmt::Debug> Visitor<C> for ElementVisitor {
+            fn visit_element(&mut self, node: &mut rstml::node::NodeElement<C>) -> bool {
+                self.elements.push(node.open_tag.name.to_string());
+                true
+            }
+        }
+        impl VisitMut for ElementVisitor {}
+
+        let visitor = ElementVisitor {
+            elements: Vec::new(),
+        };
+        let elements =
+            visit_nodes_with_custom::<_, _, EscapeCodeWalker>(&mut actual, visitor).elements;
+
+        assert_eq!(&elements, &["div"]);
+    }
+
+    #[test]
+    fn if_node_visitor_syn_expr() {
+        let tokens = quote! {
+            @if just && an || expression {
+                <div/>
+            }
+        };
+
+        let mut actual = parse_universal(tokens).into_result().unwrap();
+
+        struct ElementVisitor {
+            exprs: Vec<String>,
+        }
+        impl<C: std::fmt::Debug> Visitor<C> for ElementVisitor {}
+        impl VisitMut for ElementVisitor {
+            fn visit_expr_mut(&mut self, exprs: &mut syn::Expr) {
+                self.exprs.push(exprs.to_token_stream().to_string());
+            }
+        }
+
+        let visitor = ElementVisitor { exprs: Vec::new() };
+        let elements =
+            visit_nodes_with_custom::<_, _, EscapeCodeWalker>(&mut actual, visitor).exprs;
+
+        assert_eq!(&elements, &["just && an || expression"]);
+    }
 }
diff --git a/rstml-controll-flow/src/extendable.rs b/rstml-control-flow/src/extendable.rs
similarity index 85%
rename from rstml-controll-flow/src/extendable.rs
rename to rstml-control-flow/src/extendable.rs
index 38fafd2..2212c0c 100644
--- a/rstml-controll-flow/src/extendable.rs
+++ b/rstml-control-flow/src/extendable.rs
@@ -131,12 +131,34 @@ impl_tuple!(A, B, C, D, E, F, G);
 impl_tuple!(A, B, C, D, E, F, G, H);
 
 fn init_extendable_node<E: Tuple + 'static>() {
+    assert_context_empty();
     TO_TOKENS.with_borrow_mut(|f| *f = Some(Box::new(E::to_tokens)));
     PARSE_RECOVERABLE.with_borrow_mut(|f| *f = Some(Box::new(E::parse_recoverable)));
     PEEK.with_borrow_mut(|f| *f = Some(Box::new(E::peek)));
 }
 
-fn clear_context() {
+fn assert_context_empty() {
+    TO_TOKENS.with_borrow(|f| {
+        assert!(
+            f.is_none(),
+            "Cannot init ExtendableCustomNode context multiple times"
+        )
+    });
+    PARSE_RECOVERABLE.with_borrow(|f| {
+        assert!(
+            f.is_none(),
+            "Cannot init ExtendableCustomNode context multiple times"
+        )
+    });
+    PEEK.with_borrow(|f| {
+        assert!(
+            f.is_none(),
+            "Cannot init ExtendableCustomNode context multiple times"
+        )
+    });
+}
+
+pub fn clear_context() {
     TO_TOKENS.with_borrow_mut(|f| *f = None);
     PARSE_RECOVERABLE.with_borrow_mut(|f| *f = None);
     PEEK.with_borrow_mut(|f| *f = None);
@@ -152,6 +174,13 @@ impl ExtendableCustomNode {
         self.value.downcast_ref::<T>()
     }
 
+    /// Parses token stream into `Vec<Node<ExtendableCustomNode>>`
+    ///
+    /// Note: This function are using context from thread local storage,
+    /// after call it lefts context initiliazed, so to_tokens implementation can
+    /// be used. But second call to parse2_with_config will fail, because
+    /// context is already initialized.
+    /// You can use `clear_context` to clear context manually.
     pub fn parse2_with_config<E: Tuple + 'static>(
         config: ParserConfig,
         tokens: proc_macro2::TokenStream,
@@ -159,7 +188,6 @@ impl ExtendableCustomNode {
         init_extendable_node::<E>();
         let result =
             Parser::new(config.custom_node::<ExtendableCustomNode>()).parse_recoverable(tokens);
-        clear_context();
         result
     }
 }
diff --git a/rstml-controll-flow/src/lib.rs b/rstml-control-flow/src/lib.rs
similarity index 100%
rename from rstml-controll-flow/src/lib.rs
rename to rstml-control-flow/src/lib.rs
diff --git a/rstml-controll-flow/src/tags.rs b/rstml-control-flow/src/tags.rs
similarity index 100%
rename from rstml-controll-flow/src/tags.rs
rename to rstml-control-flow/src/tags.rs
diff --git a/src/node/mod.rs b/src/node/mod.rs
index 89738ff..f589302 100644
--- a/src/node/mod.rs
+++ b/src/node/mod.rs
@@ -268,7 +268,7 @@ pub trait CustomNode: ParseRecoverable + ToTokens {
 }
 
 /// Newtype for `std::convert::Infallible` used to implement
-/// `ToTokens`` for `Infallible``.
+/// `ToTokens` for `Infallible`.
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 pub struct Infallible(convert::Infallible);
 
diff --git a/src/visitor.rs b/src/visitor.rs
index 3fd906a..0c743aa 100644
--- a/src/visitor.rs
+++ b/src/visitor.rs
@@ -21,10 +21,12 @@ pub enum RustCode<'a> {
 /// Each method returns a bool that indicates if the visitor should continue to
 /// traverse the tree. If the method returns false, the visitor will stop
 /// traversing the tree.
-pub trait Visitor {
-    type Custom: CustomNode;
+///
+/// By default Visitor are abstract over CustomNode, but it is possible to
+/// implement a Visitor for concrete CustomNode.
+pub trait Visitor<Custom> {
     // Visit node types
-    fn visit_node(&mut self, _node: &mut Node<Self::Custom>) -> bool {
+    fn visit_node(&mut self, _node: &mut Node<Custom>) -> bool {
         true
     }
     fn visit_block(&mut self, _node: &mut NodeBlock) -> bool {
@@ -36,19 +38,19 @@ pub trait Visitor {
     fn visit_doctype(&mut self, _node: &mut NodeDoctype) -> bool {
         true
     }
-    fn visit_raw_node<OtherC: CustomNode>(&mut self, _node: &mut RawText<OtherC>) -> bool {
+    fn visit_raw_node<AnyC: CustomNode>(&mut self, _node: &mut RawText<AnyC>) -> bool {
         true
     }
-    fn visit_custom(&mut self, _node: &mut Self::Custom) -> bool {
+    fn visit_custom(&mut self, _node: &mut Custom) -> bool {
         true
     }
     fn visit_text_node(&mut self, _node: &mut NodeText) -> bool {
         true
     }
-    fn visit_element(&mut self, _node: &mut NodeElement<Self::Custom>) -> bool {
+    fn visit_element(&mut self, _node: &mut NodeElement<Custom>) -> bool {
         true
     }
-    fn visit_fragment(&mut self, _node: &mut NodeFragment<Self::Custom>) -> bool {
+    fn visit_fragment(&mut self, _node: &mut NodeFragment<Custom>) -> bool {
         true
     }
 
@@ -91,7 +93,32 @@ pub trait Visitor {
     }
 }
 
-macro_rules! visit_child {
+#[derive(Debug, Default, Clone, PartialEq, PartialOrd, Ord, Copy, Eq)]
+pub struct AnyWalker<C>(PhantomData<C>);
+
+/// Define walker for `CustomNode`.
+pub trait CustomNodeWalker {
+    type Custom: CustomNode;
+    fn walk_custom_node_fields<VisitorImpl: Visitor<Self::Custom>>(
+        visitor: &mut VisitorImpl,
+        node: &mut Self::Custom,
+    ) -> bool;
+}
+
+impl<C> CustomNodeWalker for AnyWalker<C>
+where
+    C: CustomNode,
+{
+    type Custom = C;
+    fn walk_custom_node_fields<VisitorImpl: Visitor<Self::Custom>>(
+        _visitor: &mut VisitorImpl,
+        _node: &mut C,
+    ) -> bool {
+        true
+    }
+}
+
+macro_rules! visit_inner {
     ($self:ident.$visitor:ident.$method:ident($($tokens:tt)*)) => {
         if !$self.$visitor.$method($($tokens)*) {
             return false;
@@ -112,7 +139,7 @@ macro_rules! try_visit {
 /// `syn::visit_mut::VisitMut`.
 ///
 /// For regular usecases it is recommended to use `visit_nodes`,
-/// `visit_nodes_with_custom_handler` or `visit_attributes` functions.
+/// `visit_nodes_with_custom` or `visit_attributes` functions.
 ///
 /// But if you need it can be used by calling `visit_*` methods directly.
 ///
@@ -121,14 +148,13 @@ macro_rules! try_visit {
 /// use quote::quote;
 /// use rstml::{
 ///     node::{Node, NodeText},
-///     visitor::{Visitor, VisitorWithDefault},
+///     visitor::{Visitor, Walker},
 ///     Infallible,
 /// };
 /// use syn::parse_quote;
 ///
 /// struct TestVisitor;
-/// impl Visitor for TestVisitor {
-///     type Custom = Infallible;
+/// impl<C> Visitor<C> for TestVisitor {
 ///     fn visit_text_node(&mut self, node: &mut NodeText) -> bool {
 ///         *node = parse_quote!("modified");
 ///         true
@@ -136,7 +162,7 @@ macro_rules! try_visit {
 /// }
 /// impl syn::visit_mut::VisitMut for TestVisitor {}
 ///
-/// let mut visitor = VisitorWithDefault::new(TestVisitor);
+/// let mut visitor = Walker::new(TestVisitor);
 ///
 /// let tokens = quote! {
 ///     <div>
@@ -162,41 +188,63 @@ macro_rules! try_visit {
 ///     .to_string()
 /// );
 /// ```
-pub struct VisitorWithDefault<C: CustomNode, V: Visitor + syn::visit_mut::VisitMut> {
-    custom: PhantomData<C>,
+pub struct Walker<V, C = Infallible, CW = AnyWalker<C>>
+where
+    C: CustomNode,
+    V: Visitor<C> + syn::visit_mut::VisitMut,
+    CW: CustomNodeWalker<Custom = C>,
+{
     visitor: V,
-    custom_handler: Option<Box<dyn FnMut(&mut C) -> bool>>,
+    // we use callbakc instead of marker for `CustomNodeWalker`
+    // because it will fail to resolve with infinite recursion
+    walker: PhantomData<CW>,
+    _pd: PhantomData<C>,
 }
-impl<C, V> VisitorWithDefault<C, V>
+
+impl<V, C> Walker<V, C>
 where
     C: CustomNode,
-    V: Visitor<Custom = C> + syn::visit_mut::VisitMut,
+    V: Visitor<C> + syn::visit_mut::VisitMut,
 {
     pub fn new(visitor: V) -> Self {
         Self {
-            custom: PhantomData,
             visitor,
-            custom_handler: None,
+            walker: PhantomData,
+            _pd: PhantomData,
         }
     }
-    pub fn with_custom_handler(visitor: V, handler: Box<dyn FnMut(&mut C) -> bool>) -> Self {
-        Self {
-            custom: PhantomData,
+    pub fn with_custom_handler<OtherCW>(visitor: V) -> Walker<V, C, OtherCW>
+    where
+        OtherCW: CustomNodeWalker<Custom = C>,
+    {
+        Walker {
             visitor,
-            custom_handler: Some(handler),
+            walker: PhantomData,
+            _pd: PhantomData,
         }
     }
 }
+impl<V, C, CW> Walker<V, C, CW>
+where
+    C: CustomNode,
+    V: Visitor<C> + syn::visit_mut::VisitMut,
+    CW: CustomNodeWalker<Custom = C>,
+{
+    pub fn destruct(self) -> V {
+        self.visitor
+    }
+}
 
-impl<C, V> Visitor for VisitorWithDefault<C, V>
+impl<V, C, CW> Visitor<C> for Walker<V, C, CW>
 where
     C: CustomNode,
-    V: Visitor<Custom = C> + syn::visit_mut::VisitMut,
+    V: Visitor<C> + syn::visit_mut::VisitMut,
+    CW: CustomNodeWalker<Custom = C>,
 {
-    type Custom = C;
-    fn visit_node(&mut self, node: &mut Node<Self::Custom>) -> bool {
-        visit_child!(self.visitor.visit_node(node));
+    fn visit_node(&mut self, node: &mut Node<C>) -> bool {
+        visit_inner!(self.visitor.visit_node(node));
 
+        println!("visit node");
         match node {
             Node::Block(b) => self.visit_block(b),
             Node::Comment(c) => self.visit_comment(c),
@@ -209,7 +257,7 @@ where
         }
     }
     fn visit_block(&mut self, node: &mut NodeBlock) -> bool {
-        visit_child!(self.visitor.visit_block(node));
+        visit_inner!(self.visitor.visit_block(node));
 
         match node {
             NodeBlock::Invalid(b) => self.visit_invalid_block(b),
@@ -217,36 +265,33 @@ where
         }
     }
     fn visit_comment(&mut self, node: &mut NodeComment) -> bool {
-        visit_child!(self.visitor.visit_comment(node));
+        visit_inner!(self.visitor.visit_comment(node));
 
         self.visit_rust_code(RustCode::LitStr(&mut node.value))
     }
     fn visit_doctype(&mut self, node: &mut NodeDoctype) -> bool {
-        visit_child!(self.visitor.visit_doctype(node));
+        visit_inner!(self.visitor.visit_doctype(node));
 
         self.visit_raw_node(&mut node.value)
     }
     fn visit_raw_node<OtherC: CustomNode>(&mut self, node: &mut RawText<OtherC>) -> bool {
-        visit_child!(self.visitor.visit_raw_node(node));
+        visit_inner!(self.visitor.visit_raw_node(node));
 
         true
     }
-    fn visit_custom(&mut self, node: &mut Self::Custom) -> bool {
-        visit_child!(self.visitor.visit_custom(node));
+    fn visit_custom(&mut self, node: &mut C) -> bool {
+        visit_inner!(self.visitor.visit_custom(node));
 
-        if let Some(ref mut handler) = &mut self.custom_handler {
-            handler(node)
-        } else {
-            true
-        }
+        CW::walk_custom_node_fields(self, node)
     }
     fn visit_text_node(&mut self, node: &mut NodeText) -> bool {
-        visit_child!(self.visitor.visit_text_node(node));
+        visit_inner!(self.visitor.visit_text_node(node));
 
         self.visit_rust_code(RustCode::LitStr(&mut node.value))
     }
-    fn visit_element(&mut self, node: &mut NodeElement<Self::Custom>) -> bool {
-        visit_child!(self.visitor.visit_element(node));
+    fn visit_element(&mut self, node: &mut NodeElement<C>) -> bool {
+        println!("Visit element Walker: {:?}", node.open_tag);
+        visit_inner!(self.visitor.visit_element(node));
 
         try_visit!(self.visit_open_tag(&mut node.open_tag));
 
@@ -262,8 +307,8 @@ where
         }
         true
     }
-    fn visit_fragment(&mut self, node: &mut NodeFragment<Self::Custom>) -> bool {
-        visit_child!(self.visitor.visit_fragment(node));
+    fn visit_fragment(&mut self, node: &mut NodeFragment<C>) -> bool {
+        visit_inner!(self.visitor.visit_fragment(node));
 
         for child in node.children_mut() {
             try_visit!(self.visit_node(child))
@@ -272,14 +317,14 @@ where
     }
 
     fn visit_open_tag(&mut self, open_tag: &mut OpenTag) -> bool {
-        visit_child!(self.visitor.visit_open_tag(open_tag));
+        visit_inner!(self.visitor.visit_open_tag(open_tag));
 
         try_visit!(self.visit_node_name(&mut open_tag.name));
 
         true
     }
     fn visit_close_tag(&mut self, closed_tag: &mut CloseTag) -> bool {
-        visit_child!(self.visitor.visit_close_tag(closed_tag));
+        visit_inner!(self.visitor.visit_close_tag(closed_tag));
 
         try_visit!(self.visit_node_name(&mut closed_tag.name));
 
@@ -287,7 +332,7 @@ where
     }
 
     fn visit_attribute(&mut self, attribute: &mut NodeAttribute) -> bool {
-        visit_child!(self.visitor.visit_attribute(attribute));
+        visit_inner!(self.visitor.visit_attribute(attribute));
 
         match attribute {
             NodeAttribute::Attribute(a) => self.visit_keyed_attribute(a),
@@ -295,7 +340,7 @@ where
         }
     }
     fn visit_keyed_attribute(&mut self, attribute: &mut KeyedAttribute) -> bool {
-        visit_child!(self.visitor.visit_keyed_attribute(attribute));
+        visit_inner!(self.visitor.visit_keyed_attribute(attribute));
 
         match &mut attribute.possible_value {
             KeyedAttributeValue::None => self.visit_attribute_flag(&mut attribute.key),
@@ -304,11 +349,11 @@ where
         }
     }
     fn visit_attribute_flag(&mut self, key: &mut NodeName) -> bool {
-        visit_child!(self.visitor.visit_attribute_flag(key));
+        visit_inner!(self.visitor.visit_attribute_flag(key));
         true
     }
     fn visit_attribute_binding(&mut self, key: &mut NodeName, value: &mut FnBinding) -> bool {
-        visit_child!(self.visitor.visit_attribute_binding(key, value));
+        visit_inner!(self.visitor.visit_attribute_binding(key, value));
 
         for input in value.inputs.iter_mut() {
             try_visit!(self.visit_rust_code(RustCode::Pat(input)))
@@ -320,19 +365,19 @@ where
         key: &mut NodeName,
         value: &mut AttributeValueExpr,
     ) -> bool {
-        visit_child!(self.visitor.visit_attribute_value(key, value));
+        visit_inner!(self.visitor.visit_attribute_value(key, value));
 
         self.visit_node_name(key);
         self.visit_rust_code(RustCode::Expr(&mut value.value))
     }
 
     fn visit_invalid_block(&mut self, block: &mut InvalidBlock) -> bool {
-        visit_child!(self.visitor.visit_invalid_block(block));
+        visit_inner!(self.visitor.visit_invalid_block(block));
 
         true
     }
     fn visit_node_name(&mut self, name: &mut NodeName) -> bool {
-        visit_child!(self.visitor.visit_node_name(name));
+        visit_inner!(self.visitor.visit_node_name(name));
 
         true
     }
@@ -345,7 +390,7 @@ where
                 RustCode::LitStr(l) => RustCode::LitStr(l),
                 RustCode::Pat(p) => RustCode::Pat(p),
             };
-            visit_child!(self.visitor.visit_rust_code(rewrap));
+            visit_inner!(self.visitor.visit_rust_code(rewrap));
         }
 
         match code {
@@ -365,14 +410,10 @@ where
 /// Return modified visitor back
 pub fn visit_nodes<V, C>(nodes: &mut [Node<C>], visitor: V) -> V
 where
-    V: Visitor<Custom = C> + syn::visit_mut::VisitMut,
-    <V as Visitor>::Custom: CustomNode,
+    C: CustomNode,
+    V: Visitor<C> + syn::visit_mut::VisitMut,
 {
-    let mut visitor = VisitorWithDefault {
-        custom: PhantomData,
-        visitor,
-        custom_handler: None,
-    };
+    let mut visitor = Walker::<V, C>::new(visitor);
     for node in nodes {
         visitor.visit_node(node);
     }
@@ -387,20 +428,13 @@ where
 /// and call visitor methods for its children.
 ///
 /// Return modified visitor back
-pub fn visit_nodes_with_custom_handler<V, C>(
-    nodes: &mut [Node<C>],
-    handler: Box<dyn FnMut(&mut C) -> bool>,
-    visitor: V,
-) -> V
+pub fn visit_nodes_with_custom<V, C, CW>(nodes: &mut [Node<C>], visitor: V) -> V
 where
-    V: Visitor<Custom = C> + syn::visit_mut::VisitMut,
-    <V as Visitor>::Custom: CustomNode,
+    C: CustomNode,
+    V: Visitor<C> + syn::visit_mut::VisitMut,
+    CW: CustomNodeWalker<Custom = C>,
 {
-    let mut visitor = VisitorWithDefault {
-        custom: PhantomData,
-        visitor,
-        custom_handler: Some(handler),
-    };
+    let mut visitor = Walker::with_custom_handler::<CW>(visitor);
     for node in nodes {
         visitor.visit_node(node);
     }
@@ -410,13 +444,10 @@ where
 /// Visit attributes in array calling visitor methods.
 pub fn visit_attributes<V>(attributes: &mut [NodeAttribute], visitor: V) -> V
 where
-    V: Visitor<Custom = Infallible> + syn::visit_mut::VisitMut,
+    V: Visitor<Infallible> + syn::visit_mut::VisitMut,
+    Walker<V>: Visitor<Infallible>,
 {
-    let mut visitor = VisitorWithDefault {
-        custom: PhantomData,
-        visitor,
-        custom_handler: None,
-    };
+    let mut visitor = Walker::new(visitor);
     for attribute in attributes {
         visitor.visit_attribute(attribute);
     }
@@ -436,8 +467,7 @@ mod tests {
         struct TestVisitor {
             collected_names: Vec<NodeName>,
         }
-        impl Visitor for TestVisitor {
-            type Custom = Infallible;
+        impl<C: CustomNode> Visitor<C> for TestVisitor {
             fn visit_node_name(&mut self, name: &mut NodeName) -> bool {
                 self.collected_names.push(name.clone());
                 true
@@ -469,6 +499,41 @@ mod tests {
         );
     }
 
+    #[test]
+    fn collect_node_elements() {
+        #[derive(Default)]
+        struct TestVisitor {
+            collected_names: Vec<NodeName>,
+        }
+        impl<C: CustomNode> Visitor<C> for TestVisitor {
+            fn visit_element(&mut self, node: &mut NodeElement<C>) -> bool {
+                self.collected_names.push(node.open_tag.name.clone());
+                true
+            }
+        }
+        // empty impl
+        impl syn::visit_mut::VisitMut for TestVisitor {}
+
+        let stream = quote! {
+            <div>
+                <span></span>
+                <span></span>
+            </div>
+            <!-- "comment" -->
+            <foo attr key=value> </foo>
+        };
+        let mut nodes = crate::parse2(stream).unwrap();
+        let visitor = visit_nodes(&mut nodes, TestVisitor::default());
+        // convert node_names to string;
+        let node_names = visitor
+            .collected_names
+            .iter()
+            .map(|name| name.to_string())
+            .collect::<Vec<_>>();
+
+        assert_eq!(node_names, vec!["div", "span", "span", "foo"]);
+    }
+
     #[test]
     fn collect_rust_blocks() {
         #[derive(Default)]
@@ -476,9 +541,7 @@ mod tests {
             collected_blocks: Vec<syn::Block>,
         }
         // empty impl
-        impl Visitor for TestVisitor {
-            type Custom = Infallible;
-        }
+        impl<C: CustomNode> Visitor<C> for TestVisitor {}
         impl syn::visit_mut::VisitMut for TestVisitor {
             fn visit_block_mut(&mut self, i: &mut syn::Block) {
                 self.collected_blocks.push(i.clone());
@@ -518,10 +581,8 @@ mod tests {
         struct TestVisitor {
             collected_raw_text: Vec<RawText<Infallible>>,
         }
-        impl Visitor for TestVisitor {
-            type Custom = Infallible;
-
-            fn visit_raw_node<OtherC: CustomNode>(&mut self, node: &mut RawText<OtherC>) -> bool {
+        impl<C: CustomNode> Visitor<C> for TestVisitor {
+            fn visit_raw_node<AnyC: CustomNode>(&mut self, node: &mut RawText<AnyC>) -> bool {
                 let raw = node.clone().convert_custom::<Infallible>();
                 self.collected_raw_text.push(raw);
                 true
@@ -560,9 +621,7 @@ mod tests {
         struct TestVisitor {
             collected_literals: Vec<syn::LitStr>,
         }
-        impl Visitor for TestVisitor {
-            type Custom = Infallible;
-        }
+        impl<C: CustomNode> Visitor<C> for TestVisitor {}
         impl syn::visit_mut::VisitMut for TestVisitor {
             fn visit_lit_str_mut(&mut self, i: &mut syn::LitStr) {
                 self.collected_literals.push(i.clone());
@@ -596,8 +655,7 @@ mod tests {
     #[test]
     fn modify_text_visitor() {
         struct TestVisitor;
-        impl Visitor for TestVisitor {
-            type Custom = Infallible;
+        impl<C: CustomNode> Visitor<C> for TestVisitor {
             fn visit_text_node(&mut self, node: &mut NodeText) -> bool {
                 *node = parse_quote!("modified");
                 true
@@ -605,7 +663,7 @@ mod tests {
         }
         impl syn::visit_mut::VisitMut for TestVisitor {}
 
-        let mut visitor = VisitorWithDefault::new(TestVisitor);
+        let mut visitor = Walker::new(TestVisitor);
 
         let tokens = quote! {
             <div>