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
}

/**
 Calculates the quality, angle and distance of a scan from the raw Data.
 
 - returns:
 LidarScan containing the quality, angle in degrees, and distance in mm.
 
 - throws:
 RPLidarError.SCAN_ERROR if the raw input is malformed or empty.
 
 - parameters:
 - raw: The bytes containing the scan data.
 
 */
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((Int(raw[1]) >> 1) + (Int(raw[2]) << 7)) / 64
    let distance = Float(Int(raw[3]) + (Int(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
}


/**
 Function alias for measurement (single scan) closure. Use with iterMeasurements
 
 - parameters:
 - scan: Single lidar scan.
 */
public typealias MeasurementHandler = (_ scan: LidarScan) -> Bool

/**
 Function alias for scan (multiple scans) closure. Use with iterScans
 
 - parameters:
 - scans: Array of lidar scans for this batch.
 */
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
    
    private(set) var isScanning = false
    private(set) var dataSize: UInt8 = 5
    /**
     A class to interact with the RPLidar A1 and A2.
     
     - parameters:
     onPort: Serial Port the lidar is attached to
     measurementDelayus: Delay that should be applied between each measurement. This is needed in case the serial port timeout is not working correctly.
     
     - throws If the serial port cannot be connected, or the motor cannot be started.
     */
    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.")
        }
    }
    
    /**
     Connect and configure the serial port.
     */
    public func connect() throws {
        disconnect()
        try serialPort.openPort()
        serialPort.setBaudrate(baudrate: 115200)
    }
    
    /**
     Close the serial port.
     */
    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
    }
    
    /**
     Gets information about the connected lidar.
     
     - 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)
    }
    
    /**
     Get the current health status of the lidar.
     
     - returns (Health status, errorCode
     - throws If the received health status is in the incorrect format.
     */
    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)
    }
    
    /**
     Clear all data in the serial port buffer.
     - throws If the buffer could not be emptied.
     */
    public func clearInput() throws{
        _ = try serialPort.readData(ofLength: serialPort.inWaiting)
    }
    
    /**
     Send the stop command to the lidar.
     - throws If the stop command cannot be sent.
     */
    public func stop() throws{
        try sendCommand(Constants.STOP.asData())
        isScanning = false
    }
    
    /**
     Send the reset command to the lidar. This puts the lidar into a state as though it just connected and started scannning.
     - throws If the reset command annot be reset.
     */
    public func reset() throws{
        try sendCommand(Constants.RESET.asData())
    }
    
    /**
     Iterate over the lidar measurments, calling the given measurement handler each time scan is processed.
     
     - parameters:
     maxBufferMeasurements: Max number of measurements to include in the buffer before flushing it.
     onMeasure: Measurement handler to execute on every read.
     
     - throws If the input data is in the incorrect format, or the health retrieved errors.
     */
    public func iterMeasurements(maxBufferMeasurements: Int = 500, _ onMeasure: MeasurementHandler) throws {
        if !isScanning {
            throw RPLidarError.SCAN_ERROR(message: "You must start scanning!")
        }
        var read = true
        while read {
            do {
                let raw = try receiveMeasurementAndClearBuffer(dataSize: dataSize, maxBufferMeasurements: maxBufferMeasurements)
                if raw.count > 0 {
                    read = try onMeasure(processScan(raw: raw))
                }
                if measurementDelayus > 0 {
                    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)")
            }
        }
    }
    
    /**
     Send the command to the lidar to start scanning. This should be called before using any iterator.
     */
    public func startScanning() throws -> UInt8  {
        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
        }
        isScanning = true
        self.dataSize = dataSize
        return dataSize
    }
    
    /**
     Receive a measurement from the lidar. If the data in waiting exceeds the buffer, clear it.
     - parameters:
     dataSize: The length of the data in bytes for a single lidar measurement.
     maxBufferMeasurements: The maximum length allowed in the buffer before it should be cleared.
     */
    public func receiveMeasurementAndClearBuffer(dataSize: UInt8, maxBufferMeasurements: Int = 500) throws -> Data{
        if !isScanning {
            throw RPLidarError.SCAN_ERROR(message: "Haven't started scanning")
        }
        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))
            }
        }
        return raw
    }
    
    /**
     Iterate over batches of measurements.
     - parameters:
     maxBufferMeasurements: Max number of measurements to include in the buffer before flushing it.
     minLength: Minimum number of measurements to include in a scan.
     onScan: Handler function to execute on each scan.
     
     - throws If the input data is in the incorrect format, or the health retrieved errors.
     */
    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.removeAll()
            }
            if measurement.quality > 0 && measurement.distance > 0 {
                scan.append(measurement)
            }
            return read
        }
    }
    
    /**
     Make an iterator that can be used to iterate over batches of measurements.
     - parameters:
     withLength: Minimum number of measurements to process in an iteration.
     maxBufferMeasurements: The maximum number of measurements allowed in the buffer before clearing it.
     */
    public func makeScanIterator(withLength minLength: Int = 5, maxBufferMeasurements: Int = 500) -> LidarScanIterator {
        return LidarScanIterator(lidar: self, measurementsPerBatch: minLength, maxBufferMeasurements: maxBufferMeasurements)
    }
    
    private func setPwm(_ pwm: Int) throws{
        assert(0 <= pwm && pwm <= Constants.MAX_MOTOR_PWM)
        try self.sendPayloadCommand(Constants.SET_PWM_BYTE, payload: pwm.asData())
    }
    
    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])
    }
}

// MARK: Iterator Extensions

extension SwiftRPLidar: Sequence {
    public typealias Element = LidarScan
    public typealias Iterator = LidarMeasurementIterator
    
    public func makeIterator() -> LidarMeasurementIterator {
        return LidarMeasurementIterator(lidar: self, maxBufferMeasurements: 500)
    }
}


/**
 Iterator for individual lidar measurements.
 A scan may be nil when data is not received from serialPort, or when the lidar is not ready for scanning.
 */
public struct LidarMeasurementIterator {
    let lidar: SwiftRPLidar
    let maxBufferMeasurements: Int
}

extension LidarMeasurementIterator: IteratorProtocol{
    public typealias Element = LidarScan
    
    public mutating func next() -> LidarScan? {
        if lidar.isScanning {
            do {
                let raw = try lidar.receiveMeasurementAndClearBuffer(dataSize: lidar.dataSize, maxBufferMeasurements: maxBufferMeasurements)
                if raw.count > 0 {
                    return try processScan(raw: raw)
                }
            } catch {
                print("Failed to process lidar scan")
            }
        } else {
            print("You must start scanning!")
        }
        return nil
    }
}

/**
 Iterator for batches of lidar measurements.
 */
public struct LidarScanIterator {
    let lidar: SwiftRPLidar
    let measurementsPerBatch: Int
    let maxBufferMeasurements: Int
    public typealias Element = [LidarScan]
}

extension LidarScanIterator: Sequence {
    public typealias Iterator = LidarScanIterator
    
    public func makeIterator() -> LidarScanIterator {
        return LidarScanIterator(lidar: lidar, measurementsPerBatch: measurementsPerBatch, maxBufferMeasurements: maxBufferMeasurements)
    }
}

extension LidarScanIterator: IteratorProtocol {
    
    public mutating func next() -> [LidarScan]? {
        
        if lidar.isScanning {
            var allScans: [LidarScan] = []
            do {
                while allScans.count < measurementsPerBatch {
                    let raw = try lidar.receiveMeasurementAndClearBuffer(dataSize: lidar.dataSize, maxBufferMeasurements: maxBufferMeasurements)
                    if raw.count > 0 {
                        allScans.append(try processScan(raw: raw))
                    }
                }
                return allScans
            } catch {
                print("Failed to process scan")
            }
        }
        return nil
    }
}