SwiftUI-Tree.swift

//
//  ContentView.swift
//  SwiftUITreeDemo
//
//  Created by MC on 2020/7/9.
//

import SwiftUI

struct ContentView: View {
    var body: some View {
        DiagramSampleExample1()
    }
}

extension Tree where A == Int {
    mutating func insert(_ number: Int) {
        if number < value {
            if children.count > 0 {
                children[0].insert(number)
            } else {
                children.append(Tree(number))
            }
        } else {
            if children.count == 2 {
                children[1].insert(number)
            } else if children.count == 1, children[0].value > number {
                children[0].insert(number)
            } else {
                children.append(Tree(number))
            }
        }
    }
}

struct DiagramExample: View {
    @State private var binarytree = Tree<Int>(130, children: [
        Tree<Int>(20, children: [
            Tree(21),
            Tree(22)
        ]),
        Tree<Int>(30, children: [
            Tree(31),
            Tree(32)
        ])
    ])

    var body: some View {
        VStack {
            Diagram(tree: binarytree, node: { value in
                Text("\(value)")
                    .modifier(RoundedCircleStyle())
            })
            
            Button("随机插入") {
                withAnimation {
                    self.binarytree.insert(Int.random(in: 0...100))
                }
            }
        }

    }
}

extension CGPoint: VectorArithmetic {
    public static func - (lhs: CGPoint, rhs: CGPoint) -> CGPoint {
        CGPoint(x: lhs.x - rhs.x, y: lhs.y - rhs.y)
    }

    public mutating func scale(by rhs: Double) {
        x *= CGFloat(rhs)
        y *= CGFloat(rhs)
    }

    public var magnitudeSquared: Double {
        0
    }

    public static func + (lhs: CGPoint, rhs: CGPoint) -> CGPoint {
        CGPoint(x: lhs.x + rhs.x, y: lhs.y + rhs.y)
    }
}

struct Line: Shape {
    var from: CGPoint
    var to: CGPoint

    var animatableData: AnimatablePair<CGPoint, CGPoint> {
        get {
            AnimatablePair(from, to)
        }
        set {
            from = newValue.first
            to = newValue.second
        }
    }

    func path(in rect: CGRect) -> Path {
        var path = Path()

        path.move(to: from)
        path.addLine(to: to)

        return path
    }
}

struct Diagram<A, V: View>: View {
    let tree: Tree<A>
    let node: (A) -> V

    typealias Key = CollectDict<String, Anchor<CGPoint>>

    var body: some View {
        VStack(spacing: 10) {
            node(tree.value)
                .anchorPreference(key: Key.self, value: .center, transform: {
                    [self.tree.id: $0]
                })

            HStack(alignment: .bottom, spacing: 10) {
                ForEach(tree.children) { child in
                    Diagram(tree: child, node: self.node)
                }
            }
        }
        .backgroundPreferenceValue(Key.self) { (centers: [String: Anchor<CGPoint>]) in
            GeometryReader { proxy in
                ForEach(self.tree.children) { child in
                    Line(from: proxy[centers[self.tree.id]!],
                         to: proxy[centers[child.id]!])
                        .stroke()
                }
            }
        }
    }
}

struct CollectDict<Key: Hashable, Value>: PreferenceKey {
    static var defaultValue: [Key: Value] { [:] }
    static func reduce(value: inout [Key: Value], nextValue: () -> [Key: Value]) {
        value.merge(nextValue(), uniquingKeysWith: { $1 })
    }
}

struct RoundedCircleStyle: ViewModifier {
    func body(content: Content) -> some View {
        content
            .frame(width: 50, height: 50)
            .background(Circle().stroke(Color.primary, lineWidth: /*@START_MENU_TOKEN@*/1.0/*@END_MENU_TOKEN@*/))
            .background(Circle().foregroundColor(Color.green.opacity(0.8)))
            .padding(5)
    }
}

struct DiagramSampleExample1: View {
    let binarytree = Tree<String>("爷爷", children: [
        Tree<String>("大爷", children: [
            Tree("大侄子"),
            Tree("二侄子")
        ]),
        Tree<String>("爸爸", children: [
            Tree("大儿子"),
            Tree("小儿子")
        ])
    ])

    var body: some View {
        DiagramSample(tree: binarytree, node: { value in
            Text("\(value)")
                .modifier(RoundedCircleStyle())
        })
    }
}

struct DiagramSampleExample: View {
    let binarytree = Tree<Int>(10, children: [
        Tree<Int>(20, children: [
            Tree(21),
            Tree(22)
        ]),
        Tree<Int>(30, children: [
            Tree(31),
            Tree(32)
        ])
    ])

    var body: some View {
        DiagramSample(tree: binarytree, node: { value in
            Text("\(value)")
                .modifier(RoundedCircleStyle())
        })
    }
}

struct DiagramSample<A, V: View>: View {
    let tree: Tree<A>
    let node: (A) -> V

    var body: some View {
        VStack(spacing: 10) {
            node(tree.value)

            HStack(alignment: .bottom, spacing: 10) {
                ForEach(tree.children) { child in
                    DiagramSample(tree: child, node: self.node)
                }
            }
        }
    }
}

struct Tree<A>: Identifiable {
    let id = UUID().uuidString
    var value: A
    var children: [Tree<A>] = []

    init(_ value: A, children: [Tree<A>] = []) {
        self.value = value
        self.children = children
    }
}

struct ContentView_Previews: PreviewProvider {
    static var previews: some View {
        ContentView()
    }
}