How to manually decide next focused index path in collection view?

How can I manually decide the index path for my "next focused" cell in a collection view on tvOS?

My use case is that I have a collection view with dynamic width cells that scrolls in all directions (up/down/left/right). Sometimes when a cell in one section spans across several cells in the next one, I want to be able to manually decide what cell gets focused next.

The default behaviour seems to be that the collection view selects the cell in the middle of the previous one.

My initial idea was to implement the collectionView(_:shouldUpdateFocusIn:) delegate function, and return false + trigger a focus update, when the collection view selects the "wrong" index path.

However, for some reason, shouldUpdateFocusIn gets called several times when I return false from it and causes a visible lag.

Code Block 
func collectionView(_ collectionView: UICollectionView, shouldUpdateFocusIn context: UICollectionViewFocusUpdateContext) -> Bool {
    if let nextIndexPath = shouldFocusCell(context.focusHeading, (context.previouslyFocusedIndexPath, context.nextFocusedIndexPath)) {
        // The collection view did not select the index path of the cell we want to focus next.
        // Save the correct index path and trigger a focus update.
        lastFocusChange = (lastFocusChange.next, nextIndexPath)
        setNeedsFocusUpdate()
        // Using updateFocusIfNeeded() results in the following warning:
        // WARNING: Calling updateFocusIfNeeded while a focus update is in progress. This call will be ignored.
        return false
    }
    return true
}

Next idea was to do the same thing in collectionView(_:didUpdateFocusIn:with:), but in this case we only update the focus after it has already moved to the "wrong" cell, so it becomes apparent to the user that the focus moves from the wrong cell to the correct one.

Not ideal either.

I'm using my own subclass of UICollectionViewLayout and UICollectionView, but I don't see anything I can override to be able to manually decide what index path to focus next when navigating up/down/left/right before shouldUpdateFocusIn is called.

Is there any way I can achieve this?