Sources/TwitterApacheThrift/ThriftDecoder.swift

// Copyright 2020 Twitter, Inc.
// Licensed under the Apache License, Version 2.0
// http://www.apache.org/licenses/LICENSE-2.0
//
//  ThriftDecoder.swift
//  TwitterApacheThrift
//
//  Created on 3/24/20.
//  Copyright © 2020 Twitter, Inc. All rights reserved.
//

import Foundation

/// Provide a default implementation which calls through to `Decodable`. This
/// allows `ThriftDecodable` to use the `Decodable` implementation generated by the
/// compiler.
public extension ThriftDecodable {
    init(fromThrift decoder: ThriftDecoder) throws {
        try self.init(from: decoder)
    }
}

/// Errors thrown when the thrift is decoded.
public enum ThriftDecoderError: Error {
    /// The key of the field is not the next key in the thrift, possibly because
    /// of malformed thrift data or model discrepancy
    case keyOrderMismatch
    /// Internal buffer is uninitialized should only happen if using methods
    /// other then `decode<T>(_ type: T.Type, from data: Data)` for decoding
    case uninitializedDecodingData
    /// CodingKey is missing the value
    case codingKeyMissingIntValue(key: CodingKey)
    /// The type is not encodable, example not conforming to ThriftDecodable
    case undecodableType(type: Any)
    /// The string found while decoding is not UTF8 formatted, possibly because
    /// of malformed thrift or an encoding error.
    case nonUTF8StringData(Data)
    /// The buffer overflowed while trying to read a field, possibly because
    /// of using the wrong model type
    case readBufferOverflow
    /// The thrift data has a unsupported type value
    case unsupportedThriftType
    /// The decoder unexpectedly found nil while decoding
    case unexpectedlyFoundNil
}

///An object that decodes instances of a data type from Thrift objects.
public class ThriftDecoder: Decoder {
    public var codingPath: [CodingKey] = []

    public var userInfo: [CodingUserInfoKey : Any] = [:]

    fileprivate var binary: ThriftBinary?
    fileprivate var value: ThriftObject?

    /// The specification to be used for decoding
    public var specification: ThriftSpecification = .standard

    /// Initializes `self` with defaults.
    public init() {}

    /// Initializes `self` with defaults.
    fileprivate init(value: ThriftObject, specification: ThriftSpecification) {
        self.value = value
        self.specification = specification
    }


    /// Decodes a top-level value of the given type from the given Thrift representation.
    ///
    /// - parameter type: The type of the value to decode.
    /// - parameter data: The data to decode from.
    /// - returns: A value of the requested type.
    /// - throws: An `ThriftDecoderError` if any value throws an error during decoding.
    public func decode<T>(_ type: T.Type, from data: Data) throws -> T where T : ThriftDecodable {
        self.binary = specification == .compact ? ThriftCompactBinary(data: data) : ThriftBinary(data: data)
        let thriftType = try ThriftType.decodableType(from: type.self)
        self.value = try binary?.readThrift(type: thriftType)
        return try type.init(fromThrift: self)
    }

    public func container<Key>(keyedBy type: Key.Type) throws -> KeyedDecodingContainer<Key> where Key : CodingKey {
        return KeyedDecodingContainer(KeyedContainer<Key>(decoder: self))
    }

    public func unkeyedContainer() throws -> UnkeyedDecodingContainer {
        return UnkeyedContainer(decoder: self)
    }

    public func singleValueContainer() throws -> SingleValueDecodingContainer {
        return UnkeyedContainer(decoder: self)
    }

    fileprivate func thriftCollectionContainer() throws -> UnkeyedContainer {
        return UnkeyedContainer(decoder: self)
    }

    private class KeyedContainer<Key: CodingKey>: KeyedDecodingContainerProtocol {
        var allKeys: [Key] = []

        var decoder: ThriftDecoder

        var codingPath: [CodingKey] { return [] }

        var previousFieldId: Int = 0


        init(decoder: ThriftDecoder) {
            self.decoder = decoder
        }

        func nestedContainer<NestedKey>(keyedBy type: NestedKey.Type, forKey key: Key) throws -> KeyedDecodingContainer<NestedKey> where NestedKey : CodingKey {
            return try decoder.container(keyedBy: type)
        }

        func nestedUnkeyedContainer(forKey key: Key) throws -> UnkeyedDecodingContainer {
            return try decoder.unkeyedContainer()
        }

        func superDecoder() -> Decoder {
            return decoder
        }

        func superDecoder(forKey key: Key) throws -> Decoder {
            return decoder
        }

        func decode<T>(_ type: T.Type, forKey key: Key) throws -> T where T : Decodable {
            return try decoder.decode(type, forKey: key, previousFieldId: previousFieldId)
        }

        func contains(_ key: Key) -> Bool {
            return (try? decoder.isNull(forKey: key)) == false
        }

        func decodeNil(forKey key: Key) throws -> Bool {
            return try decoder.isNull(forKey: key)
        }
    }

    fileprivate class UnkeyedContainer: UnkeyedDecodingContainer, SingleValueDecodingContainer {
        var decoder: ThriftDecoder

        var codingPath: [CodingKey] { return [] }

        var count: Int? {
            switch decoder.value {
            case .keyedCollection(let collection):
                return collection.count
            case .unkeyedCollection(let collection):
                return collection.count
            default: return nil
            }
        }

        var currentIndex: Int = 0

        var isAtEnd: Bool { return currentIndex == count }

        init(decoder: ThriftDecoder) {
            self.decoder = decoder
        }

        func nestedContainer<NestedKey>(keyedBy type: NestedKey.Type) throws -> KeyedDecodingContainer<NestedKey> where NestedKey : CodingKey {
            return try decoder.container(keyedBy: type)
        }

        func nestedUnkeyedContainer() -> UnkeyedDecodingContainer {
            return self
        }

        func superDecoder() -> Decoder {
            return decoder
        }

        func decode<T>(_ type: T.Type) throws -> T where T : Decodable {
            let value: ThriftObject
            if case .keyedCollection(let collection) = decoder.value {
                value = collection.value[currentIndex].value
            } else if case .unkeyedCollection(let collection) = decoder.value {
                value = collection.value[currentIndex]
            } else if let decoderValue = self.decoder.value {
                return try decoder.decodeType(type: type, value: decoderValue)
            } else {
                throw ThriftDecoderError.uninitializedDecodingData
            }

            return try decoder.decodeType(type: type, value: value)
        }

        func decodeKey<T>(_ type: T.Type) throws -> T where T : Decodable {
            guard case .keyedCollection(let collection) = decoder.value else {
                throw ThriftDecoderError.keyOrderMismatch
            }
            let key = collection.value[currentIndex].key
            return try decoder.decodeType(type: type, value: key)
        }

        func moveContainer() {
            currentIndex += 1
        }

        func decodeNil() -> Bool {
            return true
        }
    }
}

extension ThriftDecoder {
    private func decode<T: Decodable, Key: CodingKey>(_ type: T.Type, forKey key: Key, previousFieldId: Int) throws -> T {
        self.codingPath.append(key)
        defer {
            self.codingPath.removeLast()
        }


        guard let keyId = key.intValue else {
            throw ThriftDecoderError.codingKeyMissingIntValue(key: key)
        }
        guard case .struct(let object) = self.value else {
            throw ThriftDecoderError.keyOrderMismatch
        }
        guard let value = object.fields[keyId] else {
            throw ThriftDecoderError.unexpectedlyFoundNil
        }

        return try decodeType(type: type, value: value.data)
    }

    private func isNull<Key: CodingKey>(forKey key: Key) throws -> Bool {
        guard let keyId = key.intValue else {
            throw ThriftDecoderError.codingKeyMissingIntValue(key: key)
        }

        guard case .struct(let object) = self.value else {
            throw ThriftDecoderError.keyOrderMismatch
        }

        return object.fields[keyId] == nil
    }

    private func decodeType<T: Decodable>(type: T.Type, value: ThriftObject) throws -> T {
        let decoded: T?
        if let thriftDecodable = type as? ThriftDecodable.Type {
            decoded = try thriftDecodable.init(fromThrift: ThriftDecoder(value: value, specification: self.specification)) as? T
        } else if case .data(let data) = value {
            let binary = self.specification == .compact ? ThriftCompactBinary(data: data) : ThriftBinary(data: data)
            switch type {
            case is Bool.Type:
                decoded = try binary.readBool() as? T
            case is Data.Type:
                decoded = data as? T
            case is Double.Type:
                decoded = try binary.readDouble() as? T
            case is UInt8.Type:
                decoded = try binary.readByte() as? T
            case is Int16.Type:
                decoded = try binary.readInt16() as? T
            case is Int32.Type:
                decoded = try binary.readInt32() as? T
            case is Int64.Type:
                decoded = try binary.readInt64() as? T
            case is String.Type:
                guard let string = String(bytes: data, encoding: .utf8) else {
                    throw ThriftDecoderError.nonUTF8StringData(data)
                }
                decoded = string as? T
            default:
                decoded = nil
            }
        } else {
            decoded = nil
        }

        guard let typedDecoded = decoded else {
            throw ThriftDecoderError.undecodableType(type: type)
        }

        return typedDecoded
    }
}


extension Dictionary : ThriftDecodable where Key: Decodable, Value: Decodable {
    public init(fromThrift decoder: ThriftDecoder) throws {
        let container = try decoder.thriftCollectionContainer()
        let size = container.count ?? 0
        self.init(minimumCapacity: container.count ?? 0)
        for _ in 0..<size {
            let key = try container.decodeKey(Key.self)
            let value = try container.decode(Value.self)
            self[key] = value
            container.moveContainer()
        }
    }
}

extension Set: ThriftDecodable where Element: Decodable {
    public init(fromThrift decoder: ThriftDecoder) throws {
        let container = try decoder.thriftCollectionContainer()
        let size = container.count ?? 0
        self.init(minimumCapacity: Int(size))
        for _ in 0..<size {
            let value = try container.decode(Element.self)
            self.insert(value)
            container.moveContainer()
        }
    }
}

extension Array: ThriftDecodable where Element: Decodable {
    public init(fromThrift decoder: ThriftDecoder) throws {
        let container = try decoder.thriftCollectionContainer()
        let size = container.count ?? 0
        self.init()
        self.reserveCapacity(Int(size))
        for _ in 0..<size {
            let value = try container.decode(Element.self)
            self.append(value)
            container.moveContainer()
        }
    }
}

extension PreencodedContainer: ThriftDecodable {

    public init(fromThrift decoder: ThriftDecoder) throws {
        try self.init(from: decoder)
    }

    public init(from decoder: Decoder) throws {
        self = .value(try T(from: decoder))
    }
}