Search code examples
swiftautomatic-ref-counting

Weak var in Swift Tree implementation


I am trying to build a tree implementation in Swift to represent a chess game.

A game consists of a sequence of moves but alternative moves for a given board position are valid. I want to traverse the tree in my GUI which is why I added methods to go to a specific node in the tree.

However I am struggling with getting the memory model right. For my task I want to keep a reference to the next node and its parent node for a given node. In my understanding these should be weak in order to not introduce retain cycles. However in doing so my implementation falls apart (because I guess i do not hold a reference to the given node?).

Can somebody enlighten me on how to change my existing implementation to make this work correctly? When I remove the weak keyword my test succeeds however I do not think this is right because again because of possible retain cycles.

I removed some of the implementation to make this more readable:

import Foundation

/// GameNode represents a node of a chess game tree
public final class GameNode {
    
    // MARK: - Properties
    
    /// The position of the node
    public let position: Position
    
    /// Uniquely identifies a node
    public let nodeId: UUID
    
    /// Is the node at the top of the tree
    public let isTopNode: Bool
    
    /// The chess move that gets from the parent node to this one
    public let move: Move?
    
    /// An optional move annotation like !!, !?, ??
    public let annotation: String?
    
    /// A comment for the move
    public let comment: String?
    
    /// The parent node
    public internal(set) weak var parent: GameNode?
    
    /// Pointer to the main variation
    public internal(set) weak var next: GameNode?
    
    /// Other possible variations from this node
    public internal(set) var variations: [GameNode]
    
    // MARK: - Init
    
    /// Creates a root node
    public init(position: Position = .initial, comment: String? = nil) {
        self.position = position
        self.nodeId = UUID()
        self.isTopNode = true
        self.move = nil
        self.annotation = nil
        self.parent = nil
        self.comment = comment
        self.next = nil
        self.variations = []
    }
    
    /// Creates a node which is the result of making a move in another node
    public init(position: Position, move: Move, parent: GameNode, annotation: String? = nil, comment: String? = nil) {
        self.position = position
        self.nodeId = UUID()
        self.isTopNode = false
        self.move = move
        self.annotation = annotation
        self.parent = parent
        self.comment = comment
        self.next = nil
        self.variations = []
    }
    
    /// Reconstructs the move sequence from the start of the game to this point
    public func reconstructMovesFromBeginning() -> [Move] {
        if parent?.isTopNode == true {
            return [move].compactMap({ $0 })
        }
        
        var moves = parent?.reconstructMovesFromBeginning() ?? []
        if let move {
            moves.append(move)
        }
        
        return moves
    }
}

public final class Game {
    public private(set) var current: GameNode
    
    public init(root: GameNode = GameNode()) {
        self.current = root
    }

    var root: GameNode? {
        var tmp: GameNode? = current
        while let currentTmp = tmp, !currentTmp.isTopNode {
            tmp = currentTmp.parent
        }
        return tmp
    }
    
    public var isAtEnd: Bool {
        current.next == nil
    }
    
    public func goBackward() {
        guard let parent = current.parent else { return }
        self.current = parent
    }
    
    public func go(to node: GameNode) {
        self.current = node
    }
    
    public func play(move: Move, comment: String? = nil, annotation: String? = nil) throws {
        let newPosition = try current.position.play(move: move)
        let newNode = GameNode(position: newPosition, move: move, parent: current, annotation: annotation, comment: comment)
        
        if !isAtEnd {
            current.next?.variations.append(newNode)
        } else {
            current.next = newNode
        }
        
        go(to: newNode)
    }
    
    public var uciPath: [String] {
        current.reconstructMovesFromBeginning().map(\.uci)
    }
}

And the test:

func testGameToUCIPath() throws {
    let game = try Game(fen: "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1")
    try game.play(move: .init(from: Squares.e2, to: Squares.e4))
    try game.play(move: .init(from: Squares.e7, to: Squares.e5))
    try game.play(move: .init(from: Squares.g1, to: Squares.f3))
    try game.play(move: .init(from: Squares.b8, to: Squares.c6))
    try game.play(move: .init(from: Squares.f1, to: Squares.b5))
    XCTAssertEqual(game.uciPath, ["e2e4", "e7e5", "g1f3", "b8c6", "f1b5"])

    game.goBackward()
    XCTAssertEqual(game.uciPath, ["e2e4", "e7e5", "g1f3", "b8c6"])
    try game.play(move: .init(from: Squares.f1, to: Squares.c4))
    XCTAssertEqual(game.uciPath, ["e2e4", "e7e5", "g1f3", "b8c6", "f1c4"])
}

Solution

  • To avoid strong reference cycles you need to have clear ownership relationship in your object hierarchy. Only owner (parent) should hold strong reference to the node(s) it owns and the back references to the parent should always be weak ones.

    When we apply that rule to your code, parent will remain weak, and next node needs to be strong.

    public final class GameNode {
        ...
        /// The parent node
        public internal(set) weak var parent: GameNode?
        
        /// Pointer to the main variation
        public internal(set) var next: GameNode?
    }
    

    But, to keep that whole hierarchy alive you also need to hold strong reference to your initial root node. Otherwise, when you reassign current node its parent node will be destroyed as there will be no strong reference holing it alive.

    Which means you will need another field in the Game class

    public final class Game {
        /// initial root node 
        public private(set) var root: GameNode
    
        public private(set) var current: GameNode
        
        public init(root: GameNode = GameNode()) {
            self.root = root
            self.current = root
        }
    

    Additionally, you have another potential problem with reference cycles in the GameNode class. That is variations array which also holds strong references to nodes. If you are storing only subsequent nodes (which is what your code comment implies) from in that array then it will be fine, but if you store that particular node or any of the previous nodes you will have strong reference cycle.