Codea's Lua UITextInputTokenizer

JAMLuaTokenizer.swift

//
//  JAMLuaTokenizer.swift
//  JamKit
//
//  Created by Sim Saens on 7/4/2024.
//  Copyright © 2024 Two Lives Left. All rights reserved.
//

import UIKit

public class JAMLuaTokenizer: JAMCodeTokenizer {
    
    weak var codeInput: JAMCodeInput?
    
    var luaIdentifierCharacterSet = CharacterSet.alphanumerics.union(CharacterSet(["_"]))
    
    public override init(textInput: any UIResponder & UITextInput) {
        super.init(textInput: textInput)
        
        self.codeInput = textInput as? JAMCodeInput
    }
    
    public override func isPosition(_ position: UITextPosition, atBoundary granularity: UITextGranularity, inDirection direction: UITextDirection) -> Bool {
        guard let codeInput,
              let position = position as? JAMTextPosition,
              let (line, _, localIndex) = codeInput.buffer.lineAndLocalIndex(at: position),
              let textDirection = Direction(direction: direction) else {
            return super.isPosition(position, atBoundary: granularity, inDirection: direction)
        }
        
        switch granularity {
        case .character:
            ()
        case .word:
            if let symbol = line.symbol(atLocation: localIndex) {
                switch textDirection {
                case .forward:
                    return !symbol.range.contains(localIndex + 1)
                case .backward:
                    return !symbol.range.contains(localIndex - 1)
                }
            }
        case .sentence:
            fallthrough
        case .line:
            switch textDirection {
            case .forward:
                if localIndex == Int(line.length) - 1 {
                    return true
                } else if currentModifierFlags.isEmpty {
                    return exactCaretPlacement
                } else {
                    return false
                }
            case .backward:
                if localIndex == 0 {
                    return true
                }
            }
        case .paragraph:
            return false
        case .document:
            switch textDirection {
            case .forward:
                if position.index == codeInput.buffer.length - 1 {
                    return true
                }
            case .backward:
                if position.index == 0 {
                    return true
                }
            }
        @unknown default:
            ()
        }
        
        return super.isPosition(position, atBoundary: granularity, inDirection: direction)
    }
    
    public override func isPosition(_ position: UITextPosition, withinTextUnit granularity: UITextGranularity, inDirection direction: UITextDirection) -> Bool {
        guard let codeInput,
              let position = position as? JAMTextPosition,
              let (line, _, localIndex) = codeInput.buffer.lineAndLocalIndex(at: position),
              let textDirection = Direction(direction: direction),
              let symbol = line.symbol(atLocation: localIndex),
              granularity == .word else {
            return super.isPosition(position, withinTextUnit: granularity, inDirection: direction)
        }
        
        return switch textDirection {
        case .forward:
            symbol.range.contains(localIndex + 1)
        case .backward:
            symbol.range.contains(localIndex - 1)
        }        
    }
    
    public override func position(from position: UITextPosition, toBoundary granularity: UITextGranularity, inDirection direction: UITextDirection) -> UITextPosition? {
        guard let codeInput,
              let position = position as? JAMTextPosition,
              let (line, lineIndex, localIndex) = codeInput.buffer.lineAndLocalIndex(at: position),
              let textDirection = Direction(direction: direction),
              let symbol = line.symbol(atLocation: localIndex) else {
            return super.position(from: position, toBoundary: granularity, inDirection: direction)
        }
        
        switch granularity {
        case .word:
            if ((symbol.type&LuaSymbolType.comment.rawValue) != 0) || ((symbol.type&LuaSymbolType.documentationComment.rawValue) != 0) || ((symbol.type&LuaSymbolType.quotedString.rawValue) != 0) {
                return super.position(from: position, toBoundary: granularity, inDirection: direction)
            }
            
            let localDestination: Int
            
            switch textDirection {
            case .forward:
                // If we are not at the end of the containing symbol
                if symbol.range.upperBound > localIndex {
                    localDestination = symbol.range.upperBound
                } else if let nextSymbol = line.firstSymbol(after: localIndex, where: { symbol in
                    !symbol.shouldIgnoreInNavigation
                }) {
                    // Otherwise we must go to the start of the next symbol
                    localDestination = nextSymbol.range.lowerBound
                } else {
                    // We're at the end of the line, let the string tokenizer figure it out
                    return super.position(from: position, toBoundary: granularity, inDirection: direction)
                }
            case .backward:
                // If we are not at the start of the containing symbol
                if symbol.range.location < localIndex {
                    localDestination = symbol.range.lowerBound
                } else if let prevSymbol = line.firstSymbol(before: localIndex, where: { symbol in
                    !symbol.shouldIgnoreInNavigation
                }) {
                    // Otherwise we must go to the end of the previous symbol
                    if prevSymbol.range.upperBound < localIndex {
                        localDestination = prevSymbol.range.upperBound
                    } else {
                        // Unless that moves us nowhere, in which case go to the start
                        localDestination = prevSymbol.range.lowerBound
                    }
                    
                } else {
                    // Fall back to text system
                    return super.position(from: position, toBoundary: granularity, inDirection: direction)
                }
            }
            
            let destination = codeInput.buffer.convert(NSRange(location: localDestination, length: 0), toGlobalRangeForLineAt: UInt(lineIndex))
            
            return JAMTextPosition(int: destination.location)

        case .sentence:
            fallthrough
        case .paragraph:
            fallthrough
        case .line:
            switch textDirection {
            case .forward:
                let localResult = NSMakeRange(Int(line.length) - 1, 0)
                let globalResult = codeInput.buffer.convert(localResult, toGlobalRangeForLineAt: UInt(lineIndex))
                
                return JAMTextPosition(int: globalResult.location)
            case .backward:
                //For code editing, we want to navigate to the start of the line but not to the start
                // of the whitespace at the beginning of the line. This searches the start of the line
                // for whitespace and navigates to the start of the first character in the line when
                // attempting to navigate backwards
                if let regex = try? NSRegularExpression(pattern: "^ *", options: []) {
                    let numSpaces = regex.rangeOfFirstMatch(in: line.line, range: NSMakeRange(0, Int(line.length))).length
                    
                    let localResult = NSMakeRange(localIndex > numSpaces ? numSpaces : 0, 0)
                    let globalResult = codeInput.buffer.convert(localResult, toGlobalRangeForLineAt: UInt(lineIndex))
                    
                    return JAMTextPosition(int: globalResult.location)
                }
            }
        default:
            return super.position(from: position, toBoundary: granularity, inDirection: direction)
        }
        
        return super.position(from: position, toBoundary: granularity, inDirection: direction)
    }
    
    public override func rangeEnclosingPosition(_ position: UITextPosition, with granularity: UITextGranularity, inDirection direction: UITextDirection) -> UITextRange? {
        guard let codeInput,
              let position = position as? JAMTextPosition,
              let (line, lineIndex, localIndex) = codeInput.buffer.lineAndLocalIndex(at: position),
              let symbol = line.symbol(atLocation: localIndex),
              granularity == .word else {
            return super.rangeEnclosingPosition(position, with: granularity, inDirection: direction)
        }
        
        if ((symbol.type&LuaSymbolType.comment.rawValue) != 0) || ((symbol.type&LuaSymbolType.documentationComment.rawValue) != 0) || ((symbol.type&LuaSymbolType.quotedString.rawValue) != 0) {
            return super.rangeEnclosingPosition(position, with: granularity, inDirection: direction)
        }
        
        let symbolRange = codeInput.buffer.convert(symbol.range, toGlobalRangeForLineAt: UInt(lineIndex))
        
        if symbolRange.location != NSNotFound {
            return JAMTextRange.indexedRange(with: symbolRange)
        } else {
            return super.rangeEnclosingPosition(position, with: granularity, inDirection: direction)
        }
    }
}

// MARK: - Helpful extensions

extension JAMLuaTokenizer {
    enum Direction {
        init?(direction: UITextDirection) {
            if direction.forward {
                self = .forward
            } else if direction.backward {
                self = .backward
            } else {
                return nil
            }
        }
        
        case forward
        case backward
    }
}

private extension UITextDirection {
    var forward: Bool {
        rawValue == UITextStorageDirection.forward.rawValue ||
        rawValue == UITextLayoutDirection.right.rawValue
    }
    
    var backward: Bool {
        rawValue == UITextStorageDirection.backward.rawValue ||
        rawValue == UITextLayoutDirection.left.rawValue
    }
}

private extension JAMEditorBuffer {
    
    func lineAndLocalIndex(at position: JAMTextPosition) -> (line: JAMEditorLine, lineIndex: Int, localIndex: Int)? {
        var localIndex: Int = 0
        let lineIndex = lineNumber(forCursorLocation: position.index, origin: nil, localIndex: &localIndex)
        
        guard lineIndex != NSNotFound else {
            return nil
        }
        
        guard let line = self[lineIndex] else {
            return nil
        }
        
        return (line, lineIndex, localIndex)
    }
    
    func firstSymbol(afterLineIndex index: Int) -> (symbol: JAMEditorSymbol, lineIndex: Int)? {
        var currentIndex = index + 1
        
        while let line = self[currentIndex], line.symbols.count == 0 {
            currentIndex += 1
        }
        
        if let symbol = self[currentIndex]?.symbols.first {
            return (symbol, currentIndex)
        }
        
        return nil
    }
}

private extension JAMEditorLine {
    func symbols(after location: Int) -> [JAMEditorSymbol] {
        symbols.filter { symbol in
            symbol.range.lowerBound > location
        }
    }
    
    func symbols(before location: Int) -> [JAMEditorSymbol] {
        symbols.filter { symbol in
            symbol.range.upperBound <= location
        }.reversed()
    }
    
    func firstSymbol(after location: Int, where predicate: (JAMEditorSymbol) -> Bool) -> JAMEditorSymbol? {
        symbols(after: location).first(where: predicate)
    }
    
    func firstSymbol(before location: Int, where predicate: (JAMEditorSymbol) -> Bool) -> JAMEditorSymbol? {
        symbols(before: location).first(where: predicate)
    }
}

extension JAMEditorSymbol {
    var shouldIgnoreInNavigation: Bool {
        Set([".", ",", ":"]).contains(symbol)
    }
}