A RelaityKit entity that draws coordinate space axes

CoordinateSpaceAxesEntity.swift

//
//  CoordinateSpaceAxesEntity.swift
//
//  Created by Drew Olbrich on 7/8/23.
//  Copyright © 2023 Lunar Skydiving LLC. All rights reserved.
//
//  MIT License
//
//  Permission is hereby granted, free of charge, to any person obtaining a copy
//  of this software and associated documentation files (the "Software"), to deal
//  in the Software without restriction, including without limitation the rights
//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//  copies of the Software, and to permit persons to whom the Software is
//  furnished to do so, subject to the following conditions:
//
//  The above copyright notice and this permission notice shall be included in all
//  copies or substantial portions of the Software.
//
//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
//  SOFTWARE.
//

import UIKit
import RealityKit

/// A RealityKit entity that draws coordinate space axes.
class CoordinateSpaceAxesEntity: Entity {
    
    init(axisLength: Float, axisRadius: Float) {
        super.init()
        
        commonInit(axisLength: axisLength, axisRadius: axisRadius)
    }
    
    required init() {
        super.init()

        commonInit(axisLength: 1, axisRadius: 0.01)
    }
    
    private func commonInit(axisLength: Float, axisRadius: Float) {
        func createCylinderEntity(with color: UIColor) -> Entity {
            let cylinderMeshResource = MeshResource.generateBox(width: axisRadius*2, height: axisLength, depth: axisRadius*2)
            
            var material = PhysicallyBasedMaterial()
            material.baseColor = .init(tint: color.blended(with: .white, weight: 0.5))
            
            let modelEntity = ModelEntity(mesh: cylinderMeshResource, materials: [material])
            
            modelEntity.position = [0, axisLength/2, 0]
            
            let pivotEntity = Entity()
            pivotEntity.addChild(modelEntity)
            
            return pivotEntity
        }
        
        let xAxisEntity = createCylinderEntity(with: .systemRed)
        xAxisEntity.orientation = simd_quatf(angle: -.pi/2, axis: [0, 0, 1])
        addChild(xAxisEntity)
        
        let yAxisEntity = createCylinderEntity(with: .systemGreen)
        addChild(yAxisEntity)
        
        let zAxisEntity = createCylinderEntity(with: .systemBlue)
        zAxisEntity.orientation = simd_quatf(angle: .pi/2, axis: [1, 0, 0])
        addChild(zAxisEntity)
    }
    
}

private extension UIColor {
    
    /// Creates a dynamic color from the receiver blended with another color.
    ///
    /// If `weight` is 0, the color of `self` is returned. If `weight` is 1, the color
    /// of `otherColor` is returned.
    ///
    /// This method works correctly with light and dark mode. A dynamic color is
    /// returned, so if the user switches to dark mode, the dark mode versions of `self`
    /// and `otherColor` will be blended together and displayed.
    func blended(with otherColor: UIColor, weight: CGFloat) -> UIColor {
        func blend(_ baseColor: UIColor, with otherColor: UIColor, weight: CGFloat) -> UIColor {
            var red: CGFloat = 0
            var green: CGFloat = 0
            var blue: CGFloat = 0
            var alpha: CGFloat = 0
            baseColor.getRed(&red, green: &green, blue: &blue, alpha: &alpha)
            
            var otherRed: CGFloat = 0
            var otherGreen: CGFloat = 0
            var otherBlue: CGFloat = 0
            var otherAlpha: CGFloat = 0
            otherColor.getRed(&otherRed, green: &otherGreen, blue: &otherBlue, alpha: &otherAlpha)
            
            func lerp(_ a: CGFloat, _ b: CGFloat, _ t: CGFloat) -> CGFloat {
                return (1 - t)*a + t*b
            }
            
            return UIColor(
                red: lerp(red, otherRed, weight),
                green: lerp(green, otherGreen, weight),
                blue: lerp(blue, otherBlue, weight),
                alpha: lerp(alpha, otherAlpha, weight)
            )
        }
        
        return UIColor(dynamicProvider: { traitCollection in
            var result: UIColor = .white
            traitCollection.performAsCurrent {
                result = blend(self, with: otherColor, weight: weight)
            }
            return result
        })
    }
        
}