Caching wrapper types

[andrewparmet]

The Java protobuf implementation cleverly caches Strings and ByteStrings when calculating message size and serializing messages. This saves a string iteration - protokt has to iterate once to find the size of a String's byte array representation, and then once to write it.

We can cache similarly and even do this for wrapper types. In fact, we can do slightly better than protobuf-java on String fields by only copying lazily and never calling CodedInputStream.readString() at all - only convert to a String when requested.

This makes OptimizedSizeofConverter obsolete.

Here's a general idea of what the implementation might look like:

class CachingReference<S : Any, T : Any>(
    @Volatile private var ref: Any,
    private val converter: Converter<S, T>
) {
    val wrapped: S
        get() =
            ref.let {
                if (converter.wrapper.java.isAssignableFrom(it::class.java)) {
                    @Suppress("UNCHECKED_CAST")
                    it as S
                } else {
                    @Suppress("UNCHECKED_CAST")
                    val converted = converter.wrap(it as T)
                    ref = converted
                    converted
                }
            }

    val unwrapped: T
        get() =
            ref.let {
                if (converter.wrapped.java.isAssignableFrom(it::class.java)) {
                    @Suppress("UNCHECKED_CAST")
                    it as T
                } else {
                    @Suppress("UNCHECKED_CAST")
                    val converted = converter.unwrap(it as S)
                    ref = converted
                    converted
                }
            }
}

class CachingModel
private constructor(
    private val _name: CachingReference<String, ByteArray>,
    private val _uuid: CachingReference<java.util.UUID, ByteArray>,
    private val _instant: CachingReference<java.time.Instant, com.toasttab.protokt.Timestamp>,
    private val _duration: CachingReference<java.time.Duration, com.toasttab.protokt.Duration>?,
    val unknown: Map<Int, Unknown>
) : KtMessage {
    // Only generated if a wrapped type is used
    private constructor(
        name: Any?,
        uuid: java.util.UUID,
        instant: Any?,
        duration: Any?,
        unknown: Map<Int, Unknown>
    ) : this(
        CachingReference(name ?: "", StringConverter),
        CachingReference(uuid, UuidConverter),
        CachingReference(
            requireNotNull(instant) {
                "instant specified nonnull with (protokt.property).non_null but was null"
            },
            InstantConverter
        ),
        duration?.let { CachingReference(it, DurationConverter) },
        unknown
    )

    override val messageSize by lazy { sizeof() }

    val name: String
        get() = _name.wrapped

    val uuid: java.util.UUID
        get() = _uuid.wrapped

    val instant: java.time.Instant
        get() = _instant.wrapped

    val duration: java.time.Duration?
        get() = _duration?.wrapped

    override fun serialize(serializer: KtMessageSerializer) {
        if (_name.unwrapped.isNotEmpty()) {
            serializer.write(Tag(10)).write(_name.unwrapped)
        }
        if (_uuid.unwrapped.isNotEmpty()) {
            serializer.write(Tag(18)).write(_uuid.unwrapped)
        }
        serializer.write(Tag(50)).write(_instant.unwrapped)
        if (_duration != null) {
            serializer.write(Tag(58)).write(_duration.unwrapped)
        }
        if (unknown.isNotEmpty()) {
            serializer.writeUnknown(unknown)
        }
    }

    private fun sizeof(): Int {
        var res = 0
        if (_name.unwrapped.isNotEmpty()) {
            res += sizeof(Tag(1)) + sizeof(_name.unwrapped)
        }
        if (_uuid.unwrapped.isNotEmpty()) {
            res += sizeof(Tag(2)) + sizeof(_uuid.unwrapped)
        }
        res += sizeof(Tag(6)) + sizeof(_instant.unwrapped)
        if (_duration != null) {
            res += sizeof(Tag(7)) + sizeof(_duration.unwrapped)
        }
        res += unknown.values.sumBy { it.sizeof() }
        return res
    }

    override fun equals(other: Any?): Boolean =
        other is CachingModel &&
            other.name == name &&
            other.uuid == uuid &&
            other.instant == instant &&
            other.duration == duration &&
            other.unknown == unknown

    override fun hashCode(): Int {
        var result = unknown.hashCode()
        result = 31 * result + name.hashCode()
        result = 31 * result + uuid.hashCode()
        result = 31 * result + instant.hashCode()
        result = 31 * result + duration.hashCode()
        return result
    }

    override fun toString(): String =
        "CachingModel(" +
            "name=$name, " +
            "uuid=$uuid, " +
            "instant=$instant, " +
            "duration=$duration, " +
            "unknown=$unknown)"

    fun copy(dsl: CachingModelDsl.() -> Unit) =
        CachingModel4 {
            name = this@CachingModel.name
            uuid = this@CachingModel.uuid
            instant = this@CachingModel.instant
            duration = this@CachingModel.duration
            unknown = this@CachingModel.unknown
            dsl()
        }

    class CachingModelDsl {
        var name = ""
        var uuid: java.util.UUID? = null
        var instant: java.time.Instant? = null
        var duration: java.time.Duration? = null
        var unknown: Map<Int, Unknown> = emptyMap()
            set(newValue) { field = copyMap(newValue) }

        fun build() =
            CachingModel(
                name,
                requireNotNull(uuid) {
                    "uuid wrapped and not specified and has no default value"
                },
                instant,
                duration,
                unknown
            )
    }

    companion object Deserializer : KtDeserializer<CachingModel>, (CachingModelDsl.() -> Unit) -> CachingModel {
        override fun deserialize(deserializer: KtMessageDeserializer): CachingModel {
            var name: ByteArray? = null
            var uuid: ByteArray? = null
            var instant: com.toasttab.protokt.Timestamp? = null
            var duration: com.toasttab.protokt.Duration? = null
            var unknown: MutableMap<Int, Unknown>? = null

            while (true) {
                when (deserializer.readTag()) {
                    0 ->
                        return CachingModel(
                            name,
                            // Has to be done eagerly in case no default is possible and it would be illegal
                            UuidConverter.wrap(uuid ?: Bytes.empty.bytes),
                            instant,
                            duration,
                            finishMap(unknown)
                        )
                    10 -> name = deserializer.readByteArray()
                    18 -> uuid = deserializer.readByteArray()
                    50 -> instant =
                        deserializer.readMessage(com.toasttab.protokt.Timestamp)
                    58 -> duration =
                        deserializer.readMessage(com.toasttab.protokt.Duration)
                    else -> unknown =
                        (unknown ?: mutableMapOf()).also {
                            processUnknown(deserializer, it)
                        }
                }
            }
        }

        override fun invoke(dsl: CachingModelDsl.() -> Unit) =
            CachingModelDsl().apply(dsl).build()
    }
}