diff --git a/node.go b/node.go index 55c487e..fe5205f 100644 --- a/node.go +++ b/node.go @@ -42,6 +42,15 @@ func (n *Node) hasChild() bool { return len(n.children) > 0 } +func (n *Node) findChildByText(text string) *Node { + for _, child := range n.children { + if text == child.name { + return child + } + } + return nil +} + func (n *Node) isDirectlyUnder(node *Node) bool { if node == nil { return false diff --git a/stack.go b/stack.go index 0d7ec53..5104ac5 100644 --- a/stack.go +++ b/stack.go @@ -39,13 +39,21 @@ func (s *stack) dfs(current *Node) { size := s.size() for i := 0; i < size; i++ { parent := s.pop() - if !current.isDirectlyUnder(parent) { - continue - } - parent.addChild(current) - current.setParent(parent) - s.push(parent).push(current) + // for same name on the same hierarchy + if child := parent.findChildByText(current.name); child != nil { + if !child.isDirectlyUnder(parent) { + continue + } + s.push(parent).push(child) + } else { + if !current.isDirectlyUnder(parent) { + continue + } + parent.addChild(current) + current.setParent(parent) + s.push(parent).push(current) + } return } } diff --git a/tree_handler_output_test.go b/tree_handler_output_test.go index 01d5832..932413c 100644 --- a/tree_handler_output_test.go +++ b/tree_handler_output_test.go @@ -195,6 +195,42 @@ a err: nil, }, }, + + { + name: "case(succeeded/same name on the same hierarchy)", + in: in{ + input: strings.NewReader(strings.TrimSpace(` +- a + - same_a + - same_b + - k + - kk + - t + - p + - q + - same_a + - o + - same_b + - ppp + - g`))}, + out: out{ + output: strings.TrimPrefix(` +a +├── same_a +│ ├── same_b +│ │ ├── k +│ │ ├── kk +│ │ └── ppp +│ ├── t +│ └── o +├── p +│ └── q +└── g +`, "\n"), + err: nil, + }, + }, + { name: "case(succeeded/very deeply)", in: in{ @@ -318,11 +354,8 @@ root dir aaa out: out{ output: strings.TrimPrefix(` parent -├── child -│ ├── chilchil -│ ├── chilchil -│ └── chilchil └── child + └── chilchil `, "\n"), err: nil, }, diff --git a/tree_handler_programmably.go b/tree_handler_programmably.go index c75e196..87d72e8 100644 --- a/tree_handler_programmably.go +++ b/tree_handler_programmably.go @@ -88,12 +88,3 @@ func (parent *Node) Add(text string) *Node { parent.addChild(current) return current } - -func (parent *Node) findChildByText(text string) *Node { - for _, child := range parent.children { - if text == child.name { - return child - } - } - return nil -}