io.wavebeans.execution.distributed.SerializableMedium.kt Maven / Gradle / Ivy

Go to download
package io.wavebeans.execution.distributed

import io.wavebeans.execution.medium.Medium
import io.wavebeans.execution.medium.MediumBuilder
import kotlinx.serialization.KSerializer
import kotlinx.serialization.Serializable
import kotlinx.serialization.SerializationException
import kotlinx.serialization.descriptors.SerialDescriptor
import kotlinx.serialization.descriptors.buildClassSerialDescriptor
import kotlinx.serialization.encoding.CompositeDecoder
import kotlinx.serialization.encoding.Decoder
import kotlinx.serialization.encoding.Encoder
import kotlin.reflect.jvm.jvmName

class SerializableMediumBuilder : MediumBuilder {
    override fun from(objects: List): Medium = SerializableMedium(objects)
}

/**
 * [Medium] for distributed execution.
 */
@Serializable(with = SerializableMediumSerializer::class)
class SerializableMedium(
        val items: List
) : Medium {

    override fun extractElement(at: Int): Any? {
        return if (at < items.size) items[at] else null
    }

    override fun equals(other: Any?): Boolean {
        if (this === other) return true
        if (javaClass != other?.javaClass) return false

        other as SerializableMedium

        if (items != other.items) return false

        return true
    }

    override fun hashCode(): Int {
        return items.hashCode()
    }
}

object SerializableMediumSerializer : KSerializer {
    override val descriptor: SerialDescriptor = buildClassSerialDescriptor(SerializableMedium::class.jvmName) {
        element("items", ListObjectSerializer.descriptor)
    }

    override fun deserialize(decoder: Decoder): SerializableMedium {
        val dec = decoder.beginStructure(descriptor)
        var l: List? = null
        loop@ while (true) {
            when (val i = dec.decodeElementIndex(descriptor)) {
                CompositeDecoder.DECODE_DONE -> break@loop
                0 -> l = dec.decodeSerializableElement(descriptor, i, ListObjectSerializer)
                else -> throw SerializationException("Unknown index $i")
            }
        }

        dec.endStructure(descriptor)
        return SerializableMedium(l!!)
    }

    override fun serialize(encoder: Encoder, value: SerializableMedium) {
        val s = encoder.beginStructure(descriptor)
        s.encodeSerializableElement(descriptor, 0, ListObjectSerializer, value.items)
        s.endStructure(descriptor)
    }

}