swiftrplidar/Sources/SwiftRPLidar/SwiftRPLidar.swift

import Foundation

struct Constants{
    static let SYNC: UInt8 = 0xA5
    static let SYNC2: UInt8 = 0x5A
    
    static let GET_INFO: UInt8 = 0x50
    static let GET_HEALTH: UInt8 = 0x52
    
    static let STOP: UInt8 = 0x25
    static let RESET: UInt8 = 0x40
    
    static let SCAN: UInt8 = 0x20
    static let FORCE_SCAN: UInt8 = 0x21
    
    static let DESCRIPTOR_LEN = 7
    static let INFO_LEN = 20
    static let HEALTH_LEN = 3
    
    static let INFO_TYPE: UInt8 = 4
    static let HEALTH_TYPE: UInt8 = 6
    static let SCAN_TYPE: UInt8 = 129
    
    static let SET_PWM_BYTE: UInt8 = 0xF0
    static let MAX_MOTOR_PWM = 1023
    static let DEFAULT_MOTOR_PWM = 660
}

public enum HEALTH_STATUSES: UInt8 {
    case GOOD = 0, WARNING, ERROR
}

public enum RPLidarError: Error {
    case FAILED_TO_START,
         INCORRECT_INFO_FORMAT,
         INCORRECT_HEALTH_FORMAT,
         SCAN_ERROR(message: String),
         INCORRECT_DESCRIPTOR_FORMAT
}

func processScan(raw: Data) throws -> LidarScan {
    let newScan = raw[0] & 0b1
    let inversedNewScan = (raw[0] >> 1) & 0b1
    let quality = raw[0] >> 2
    if newScan == inversedNewScan {
        throw RPLidarError.SCAN_ERROR(message: "New scan flags mismatch")
    }
    if (raw[1] & 0b1) != 1 {
        throw RPLidarError.SCAN_ERROR(message: "Check bit not equal to 1")
    }
    let angle = (Float(raw[1] >> 1) + Float(raw[2] << 7)) / 64
    let distance = (Float(raw[3]) + Float(raw[4] << 8)) / 4
    return LidarScan(newScan: newScan == 1, quality: quality, angle: angle, distance: distance)
}

public struct LidarScan{
    public let newScan: Bool
    public let quality: UInt8
    public let angle: Float
    public let distance: Float
}


public typealias MeasurementHandler = (_ scan: LidarScan) -> Bool
public typealias ScanHandler = (_ scans: [LidarScan]) -> Bool

public class SwiftRPLidar {
    private var motor: Bool = false
    private var serialPort: LidarSerial
    private var motorRunning = false
    private let measurementDelayus: UInt32
    
    public init(onPort serialPort: LidarSerial, measurementDelayus: UInt32 = 500) throws {
        self.serialPort = serialPort
        self.measurementDelayus = measurementDelayus
        try connect()
        try startMotor()
    }
    
    deinit {
        disconnect()
        do {
            try stop()
            try stopMotor()
        } catch {
            print("Failed to stop lidar/motor.")
        }
    }
    
    public func connect() throws {
        disconnect()
        try serialPort.openPort()
        serialPort.setBaudrate(baudrate: 115200)
    }
    
    public func disconnect(){
        // Need to close, SwiftSerial is blocking.
        serialPort.closePort()
    }
    
    public func startMotor() throws{
        // A1
        serialPort.dtr = true
        // A2
        try self.setPwm(Constants.DEFAULT_MOTOR_PWM)
        self.motorRunning = true
    }
    
    public func stopMotor() throws{
        try setPwm(0)
        serialPort.dtr = false
        motorRunning = false
    }
    
    public func setPwm(_ pwm: Int) throws{
        assert(0 <= pwm && pwm <= Constants.MAX_MOTOR_PWM)
        try self.sendPayloadCommand(Constants.SET_PWM_BYTE, payload: pwm.asData())
    }
    
    // Returns (model, firmware, hardware, serialNumber)
    public func getInfo() throws -> (UInt8, (UInt8, UInt8), UInt8, String){
        try sendCommand(Data([Constants.GET_HEALTH]))
        let (dataSize, isSingle, dataType) = try readDescriptor()!
        if dataSize != Constants.INFO_LEN || !isSingle || dataType != Constants.INFO_TYPE {
            throw RPLidarError.INCORRECT_INFO_FORMAT
        }
        
        let raw = try readResponse(dataSize)
        let serialNumber = String(bytes: raw[4...], encoding: .ascii)!
        return (raw[0], (raw[2], raw[1]), raw[3], serialNumber)
    }
    
    public func getHealth() throws -> (HEALTH_STATUSES, UInt8){
        try sendCommand(Constants.GET_HEALTH.asData())
        guard let (dataSize, isSingle, dataType) = try readDescriptor() else {
            throw RPLidarError.INCORRECT_HEALTH_FORMAT
        }
        if dataSize != Constants.HEALTH_LEN || !isSingle || dataType != Constants.HEALTH_TYPE{
            throw RPLidarError.INCORRECT_HEALTH_FORMAT
        }
        let raw = try readResponse(dataSize)
        let status = HEALTH_STATUSES(rawValue: raw[0])!
        let errorCode = raw[1] << 8 + raw[2]
        return (status, errorCode)
    }
    
    public func clearInput() throws{
        _ = try serialPort.readData(ofLength: serialPort.inWaiting)
    }
    
    public func stop() throws{
        try sendCommand(Constants.STOP.asData())
    }
    
    public func reset() throws{
        try sendCommand(Constants.RESET.asData())
    }
    
    public func iterMeasurements(maxBufferMeasurements: Int = 500, _ onMeasure: MeasurementHandler) throws {
        try startMotor()
        let (status, _) = try getHealth()
        if status == .ERROR{
            throw RPLidarError.INCORRECT_HEALTH_FORMAT
        }
        
        else if status == .WARNING {
            print("Warning given when checking health.")
        }
        
        try sendCommand(Constants.SCAN.asData())
        let (dataSize, isSingle, dataType) = try readDescriptor()!
        if dataSize != 5 || isSingle || dataType != Constants.SCAN_TYPE {
            throw RPLidarError.INCORRECT_DESCRIPTOR_FORMAT
        }
        var read = true
        while read {
            do {
                let raw = try readResponse(Int(dataSize))
                if maxBufferMeasurements > 0 {
                    let dataInWaiting = serialPort.inWaiting
                    if dataInWaiting > maxBufferMeasurements {
                        print("Too many measurements in the input buffer. Clearing Buffer")
                        _ = try serialPort.readData(ofLength: dataInWaiting / Int(dataSize) * Int(dataSize))
                    }
                }
                read = try onMeasure(processScan(raw: raw))
                // TODO: Figure out why this delay is needed.
                // If some delay is not present, then there is a high chance that the scan will be incorrect, and processScan will error.
                usleep(measurementDelayus)
            } catch RPLidarError.INCORRECT_DESCRIPTOR_FORMAT {
                print("Incorrect Descriptor, skipping scan")
            } catch RPLidarError.SCAN_ERROR(let message) {
                print("Scan processing failed, skipping scan: \(message)")
            }
        }
    }
    
    public func iterScans(maxBufferMeasurements: Int = 500, minLength: Int = 5, _ onScan: ScanHandler) throws {
        var scan: [LidarScan] = []
        var read = true
        try iterMeasurements(maxBufferMeasurements: maxBufferMeasurements){ measurement in
            if measurement.newScan {
                if scan.count > minLength {
                    read = onScan(scan)
                }
                scan = []
            }
            if measurement.quality > 0 && measurement.distance > 0 {
                scan.append(measurement)
            }
            return read
        }
    }
    
    private func sendPayloadCommand(_ cmd: Int, payload: Data) throws{
        let size = UInt8(payload.count)
        var req = Constants.SYNC.asData() + cmd.asData() + size.asData() + payload
        let checksum = calcChecksum(req)
        req += checksum.asData()
        _ = try serialPort.writeData(req)
    }
    
    private func sendPayloadCommand(_ cmd: UInt8, payload: Data) throws {
        return try sendPayloadCommand(Int(cmd), payload: payload)
    }
    
    private func calcChecksum(_ data: Data) -> Int{
        var checksum = 0;
        data.forEach { body in
            checksum ^= Int(body)
        }
        return checksum
    }
    
    private func sendCommand(_ command: Data) throws{
        var req = Constants.SYNC.asData()
        req.append(command)
        _ = try serialPort.writeData(req)
    }
    
    private func readDescriptor() throws -> (UInt8, Bool, UInt8)?{
        let descriptor = try serialPort.readData(ofLength: Constants.DESCRIPTOR_LEN)
        if (descriptor.count != Constants.DESCRIPTOR_LEN){
            return nil
        }
        else if (descriptor[0...1] != Data([Constants.SYNC, Constants.SYNC2])){
            return nil
        }
        let isSingle = descriptor[descriptor.count - 2] == 0
        return (descriptor[2], isSingle, descriptor[descriptor.count - 1])
    }
    
    private func readResponse(_ dataSize: Int) throws -> Data {
        let data = try serialPort.readData(ofLength: dataSize)
        if(data.count != dataSize){
            throw RPLidarError.INCORRECT_DESCRIPTOR_FORMAT
        }
        return data
    }
    
    private func readResponse(_ dataSize: UInt8) throws -> Data {
        return try readResponse(Int(dataSize))
    }
}

// MARK: Convenience packing extensions.

extension Int{
    func asData() -> Data {
        var int = self
        return Data(bytes: &int, count: int / Int(UINT8_MAX))
    }
}

extension UInt8 {
    func asData() -> Data {
        return Data([self])
    }
}