
//
// marshaling/unmarshaling context
//

alias MarshalSize = UInt32;

private record MarshalPointerTracker (
    id: MarshalSize,
    sized?: Bool,
    marshaled?: Bool,
);

record MarshalContext (
    pointerCount: MarshalSize,
    pointers: HashMap[RawPointer, MarshalPointerTracker],
);

overload MarshalContext() = MarshalContext(
    MarshalSize(0), HashMap[RawPointer, MarshalPointerTracker]()
);

record UnmarshalContext (
    pointers: HashMap[MarshalSize, RawPointer],
);



//
// marshaling object protocol
//

// [T] marshaledSize(context: MarshalContext, x: T) MarshalSize
// -- should not throw
define marshaledSize;
// [T] marshalTo(context: MarshalContext, buffer: Pointer[Byte], x: T) Pointer[Byte]
// -- should not throw
define marshalTo;
// [T] unmarshalFrom(
//    context: UnmarshalContext,
//    begin: Pointer[Byte],
//    end: Pointer[Byte],
//    static T
// )
//     UnmarshaledType(T), Pointer[Byte]
// -- may throw
// -- must throw UnmarshalUnexpectedEOF if unmarshal code would have to read past end
define unmarshalFrom;
// [T] UnmarshaledType(static T) T2
// -- type that value of type T should be unmarshaled into
UnmarshaledType(T) = T;

record UnmarshalUnexpectedEOF ();
instance Exception (UnmarshalUnexpectedEOF);

Marshalable?(T)
    =     CallDefined?(marshaledSize, MarshalContext, T)
      and CallDefined?(marshalTo, MarshalContext, Pointer[Byte], T)
      and CallDefined?(unmarshalFrom,
            UnmarshalContext, Pointer[Byte], Pointer[Byte], Static[T]
      );

overload Marshalable?(static MarshalContext) = false;
overload Marshalable?(static UnmarshalContext) = false;



//
// bitwise-marshalable implementation
//
BitwiseMarshalable?(T) = false;

[T | BitwiseMarshalable?(T)]
overload marshaledSize(context, x: T) = MarshalSize(TypeSize(T));

[T | BitwiseMarshalable?(T)]
overload marshalTo(context, buffer: Pointer[Byte], x: T) {
    Pointer[T](buffer)^ <-- x;
    return buffer + TypeSize(T);
}

[T | BitwiseMarshalable?(T)]
overload unmarshalFrom(context, begin: Pointer[Byte], end: Pointer[Byte], static T)
    --> returned: T, newBegin: Pointer[Byte]
{
    var newBegin_ = begin + TypeSize(T);
    if (newBegin_ > end)
        throw UnmarshalUnexpectedEOF();
    copyNonoverlappingMemory(Pointer[Byte](&returned), begin, newBegin_);
    newBegin <-- newBegin_;
}

overload BitwiseMarshalable?(static Bool) = true;
[C | Character?(C)] overload BitwiseMarshalable?(static C) = true;
[I | Integer?(I)] overload BitwiseMarshalable?(static I) = true;
[F | Float?(F)] overload BitwiseMarshalable?(static F) = true;
[E | Enum?(E)] overload BitwiseMarshalable?(static E) = true;
[R | RegularRecord?(R) and allValues?(BitwiseMarshalable?, ..RecordFieldTypes(R))]
overload BitwiseMarshalable?(static R) = true;
[..TupleMembers | allValues?(BitwiseMarshalable?, ..TupleMembers)]
overload BitwiseMarshalable?(static Tuple[..TupleMembers]) = true;
[..UnionMembers | allValues?(BitwiseMarshalable?, ..UnionMembers)]
overload BitwiseMarshalable?(static Union[..UnionMembers]) = true;

[T, n | BitwiseMarshalable?(T)] overload BitwiseMarshalable?(static Array[T, n]) = true;



//
// memberwise record marshaling
//
private RecordMarshalable?(R) = allValues?(Marshalable?, ..RecordFieldTypes(R));
private RecordMemberwiseMarshalable?(R)
    = RecordMarshalable?(R) and not BitwiseMarshalable?(R);

[R | Record?(R) and RecordMemberwiseMarshalable?(R)]
overload marshaledSize(context, r: R)
    = ..foldValues(add,
        ..mapValues(f -> marshaledSize(context, f), ..recordFields(r))
    );

marshalMemberwise(context, buffer:Pointer[Byte], forward ..fields) {
    ..for (x in fields)
        buffer = marshalTo(context, buffer, x);
}

[R | Record?(R) and RecordMemberwiseMarshalable?(R)]
overload marshalTo(context, buffer_: Pointer[Byte], forward r: R) --> buffer: Pointer[Byte] {
    buffer <-- buffer_;
    marshalMemberwise(context, buffer, ..recordFields(r));
}

[FieldT]
unmarshalMemberwise(
    context,
    begin: Pointer[Byte],
    end: Pointer[Byte],
    field: FieldT,
    ..fields
) {
    field, begin <-- ..unmarshalFrom(context, begin, end, FieldT);
    try {
        unmarshalMemberwise(context, begin, end, ..fields);
    } catch (ex) {
        destroy(field);
        throw ex;
    }
}

overload unmarshalMemberwise(context, begin: Pointer[Byte], end: Pointer[Byte]) {}

[R | Record?(R) and RecordMemberwiseMarshalable?(R)]
overload unmarshalFrom(context, begin_: Pointer[Byte], end: Pointer[Byte], static R)
    --> returned: R, begin: Pointer[Byte]
{
    begin <-- begin_;
    unmarshalMemberwise(context, begin, end, ..recordFields(returned));
}



//
// memberwise tuple marshaling
//
[..T] private TupleMarshalable?(static Tuple[..T]) = allValues?(Marshalable?, ..T);
[T] private TupleMemberwiseMarshalable?(static T)
    = TupleMarshalable?(T) and not BitwiseMarshalable?(T);

[T | Tuple?(T) and TupleMemberwiseMarshalable?(T)]
overload marshaledSize(context, t: T)
    = ..foldValues(add,
        ..mapValues(f -> marshaledSize(context, f), ..unpack(t))
    );

[T | Tuple?(T) and TupleMemberwiseMarshalable?(T)]
overload marshalTo(context, buffer_: Pointer[Byte], forward t: T) --> buffer: Pointer[Byte] {
    buffer <-- buffer_;
    ..for (x in unpack(t))
        buffer = marshalTo(context, buffer, x);
}

[T | Tuple?(T) and TupleMemberwiseMarshalable?(T)]
overload unmarshalFrom(context, begin_: Pointer[Byte], end: Pointer[Byte], static T)
    --> returned: T, begin: Pointer[Byte]
{
    begin <-- begin_;
    unmarshalMemberwise(context, begin, end, ..unpackRef(returned));
}



//
// memberwise variant marshaling
//
private VariantMarshalable?(V) = allValues?(Marshalable?, ..VariantMembers(V));
private VariantMemberwiseMarshalable?(V)
    = VariantMarshalable?(V) and not BitwiseMarshalable?(V);

[V | Variant?(V) and VariantMemberwiseMarshalable?(V)]
overload marshaledSize(context, v: V)
    = MarshalSize(TypeSize(Int)) + marshaledSize(context, *v);

[V | Variant?(V) and VariantMemberwiseMarshalable?(V)]
overload marshalTo(context, buffer_: Pointer[Byte], v: V) --> buffer: Pointer[Byte] {
    buffer <-- buffer_;
    buffer = marshalTo(context, buffer, variantTag(v));
    buffer = marshalTo(context, buffer, *v);
}

record UnmarshalInvalidVariantTag (tag: Int);
instance Exception (UnmarshalInvalidVariantTag);

[..T, n]
private unmarshalVariantMember(
    context,
    begin: Pointer[Byte],
    end: Pointer[Byte],
    repr: VariantRepr[..T],
    static n
) {
    throw UnmarshalInvalidVariantTag(repr.tag);
}

[..T, n | n < countValues(..T)]
overload unmarshalVariantMember(
    context,
    begin: Pointer[Byte],
    end: Pointer[Byte],
    repr: VariantRepr[..T],
    static n
) {
    if (repr.tag == Int(n))
        staticIndex(repr.data, static n), begin
            <-- ..unmarshalFrom(context,
                begin, end, Type(staticIndex(repr.data, static n))
            );
    else
        unmarshalVariantMember(context, begin, end, repr, static n+1);
}


[V | Variant?(V) and VariantMemberwiseMarshalable?(V)]
overload unmarshalFrom(context, begin_: Pointer[Byte], end: Pointer[Byte], static V)
    --> returned: V, begin: Pointer[Byte]
{
    begin <-- begin_;
    ref repr = variantRepr(returned);
    repr.tag, begin <-- ..unmarshalFrom(context, begin, end, Int);

    unmarshalVariantMember(context, begin, end, repr, static MarshalSize(0));
}



//
// general sequence marshaling
//
private MarshalableSequence?(S)
    = SequenceContainer?(S) and Marshalable?(SequenceElementType(S));

overload MarshalableSequence?(static StringConstant) = true;

private UnmarshalAsVectorSequence?(S)
    = MarshalableSequence?(S) and not CallDefined?(push, S, SequenceElementType(S));

private sequenceMarshaledSize(context, seq)
    = MarshalSize(TypeSize(MarshalSize))
    + reduce(add, MarshalSize(0), mapped(f -> marshaledSize(context, f), seq));

[S | MarshalableSequence?(S)]
overload marshaledSize(context, seq: S) = sequenceMarshaledSize(context, seq);

[S | MarshalableSequence?(S)]
overload marshalTo(context, buffer_: Pointer[Byte], forward seq: S) --> buffer: Pointer[Byte] {
    buffer <-- buffer_;
    var sizeBuffer = buffer;
    buffer += TypeSize(MarshalSize);
    var n = MarshalSize(0);
    for (x in seq) {
        inc(n);
        buffer = marshalTo(context, buffer, x);
    }
    marshalTo(context, sizeBuffer, n);
}

[S | MarshalableSequence?(S) and SizedSequence?(S)]
overload marshalTo(context, buffer_: Pointer[Byte], forward seq: S) --> buffer: Pointer[Byte] {
    buffer <-- marshalTo(context, buffer_, MarshalSize(size(seq)));
    for (x in seq)
        buffer = marshalTo(context, buffer, x);
}

[S | MarshalableSequence?(S)]
overload unmarshalFrom(context, begin_: Pointer[Byte], end: Pointer[Byte], static S) {
    var n, begin = ..unmarshalFrom(context, begin_, end, MarshalSize);
    var s = UnmarshaledType(S)();

    for (x in range(n)) {
        var n, begin2 = ..unmarshalFrom(context, begin, end, SequenceElementType(S));
        push(s, move(n));
        begin = begin2;
    }
    return move(s), begin;
}

[S | UnmarshalAsVectorSequence?(S)]
overload UnmarshaledType(static S) = Vector[UnmarshaledType(SequenceElementType(S))];


//
// null marshaling of static values
//
[x] overload marshaledSize(context, static x) = MarshalSize(0);
[x] overload marshalTo(context, buffer: Pointer[Byte], static x) = buffer;
[x] overload unmarshalFrom(context,
    begin: Pointer[Byte], end: Pointer[Byte], static Static[x]
) = static x, begin;



//
// null marshaling of empty records, tuples
//
[R | RegularRecord?(R) and countValues(..RecordFieldTypes(R)) == 0]
overload marshaledSize(context, r: R) = MarshalSize(0);
[R | RegularRecord?(R) and countValues(..RecordFieldTypes(R)) == 0]
overload marshalTo(context, buffer: Pointer[Byte], r: R) = buffer;
[R | RegularRecord?(R) and countValues(..RecordFieldTypes(R)) == 0]
overload unmarshalFrom(context, begin: Pointer[Byte], end: Pointer[Byte], static R)
    = R(), begin;

[T | T == Type([])]
overload marshaledSize(context, r: T) = MarshalSize(0);
[T | T == Type([])]
overload marshalTo(context, buffer: Pointer[Byte], r: T) = buffer;
[T | T == Type([])]
overload unmarshalFrom(context, begin: Pointer[Byte], end: Pointer[Byte], static T)
    = [], begin;



//
// unique pointers
//
[T]
overload marshaledSize(context, p: UniquePointer[T]) {
    if (null?(p))
        return marshaledSize(context, false);
    else
        return marshaledSize(context, true) + marshaledSize(context, p^);
}

[T]
overload marshalTo(context, buffer_: Pointer[Byte], p: UniquePointer[T]) {
    if (null?(p))
        return marshalTo(context, buffer_, false);
    else {
        var buffer = marshalTo(context, buffer_, true);
        return marshalTo(context, buffer, p^);
    }
}

[T]
overload unmarshalFrom(
    context,
    begin_: Pointer[Byte], end: Pointer[Byte],
    static UniquePointer[T]
) --> p: UniquePointer[UnmarshaledType(T)], begin: Pointer[Byte] {
    var isNotNull? = Bool();
    isNotNull?, begin <-- ..unmarshalFrom(context, begin_, end, Bool);

    if (isNotNull?) {
        p <-- UniquePointer(allocateRawMemory(UnmarshaledType(T), MarshalSize(1)));
        try {
            p^, begin <-- ..unmarshalFrom(context, begin, end, T);
        } catch (ex) {
            freeRawMemory(p.ptr);
            throw ex;
        }
    } else
        p <-- nullUnique(UnmarshaledType(T));
}

[T]
overload UnmarshaledType(static UniquePointer[T]) = UniquePointer[UnmarshaledType(T)];



//
// shared pointers
//
private alias SHARED_POINTER_NULL = 0uss;
private alias SHARED_POINTER_REF = 1uss;
private alias SHARED_POINTER_DATA = 2uss;

[T]
overload marshaledSize(context: MarshalContext, p: SharedPointer[T]) {
    if (null?(p))
        return marshaledSize(context, SHARED_POINTER_NULL);
    else {
        var rawPtr = RawPointer(p.ptr);
        var pointerState = lookup(context.pointers, rawPtr);
        if (null?(pointerState)) {
            var id = context.pointerCount;
            context.pointers[rawPtr]
                = MarshalPointerTracker(id, true, false);
            context.pointerCount += 1;
            return marshaledSize(context, SHARED_POINTER_DATA)
                + marshaledSize(context, id)
                + marshaledSize(context, p^);
        } else if (not pointerState^.sized?) {
            pointerState^.sized? = true;
            return marshaledSize(context, SHARED_POINTER_DATA)
                + marshaledSize(context, pointerState^.id)
                + marshaledSize(context, p^);
        } else {
            return marshaledSize(context, SHARED_POINTER_REF)
                + marshaledSize(context, pointerState^.id);
        }
    }
}

[T]
overload marshalTo(context: MarshalContext, buffer_: Pointer[Byte], p: SharedPointer[T]) {
    var rawPtr = RawPointer(p.ptr);
    if (null?(p))
        return marshalTo(context, buffer_, SHARED_POINTER_NULL);
    else {
        var pointerState = lookup(context.pointers, rawPtr);
        if (null?(pointerState)) {
            var id = context.pointerCount;
            context.pointers[rawPtr]
                = MarshalPointerTracker(id, false, true);
            context.pointerCount += 1;
            var buffer = marshalTo(context, buffer_, SHARED_POINTER_DATA);
            buffer = marshalTo(context, buffer, id);
            buffer = marshalTo(context, buffer, p^);
            return buffer;
        } else if (not pointerState^.marshaled?) {
            pointerState^.marshaled? = true;
            var buffer = marshalTo(context, buffer_, SHARED_POINTER_DATA);
            buffer = marshalTo(context, buffer, pointerState^.id);
            buffer = marshalTo(context, buffer, p^);
            return buffer;
        } else {
            var buffer = marshalTo(context, buffer_, SHARED_POINTER_REF);
            buffer = marshalTo(context, buffer, pointerState^.id);
            return buffer;
        }
    }
}

record UnmarshalIncompleteObjectGraph ();
record UnmarshalInconsistentObjectGraph ();
record UnmarshalInvalidObjectGraphData ();

instance Exception (
    UnmarshalIncompleteObjectGraph,
    UnmarshalInconsistentObjectGraph,
    UnmarshalInvalidObjectGraphData,
);

[T]
overload unmarshalFrom(context: UnmarshalContext,
    begin_: Pointer[Byte], end: Pointer[Byte],
    static SharedPointer[T]
) {
    var mode, begin = ..unmarshalFrom(context, begin_, end, UInt8); 
    switch (mode)
    case (SHARED_POINTER_NULL) {
        return nullShared(UnmarshaledType(T)), begin;
    }
    case (SHARED_POINTER_REF) {
        var id = MarshalSize();
        id, begin = ..unmarshalFrom(context, begin, end, MarshalSize);
        var pointerState = lookup(context.pointers, id);

        if (null?(pointerState))
            throw UnmarshalIncompleteObjectGraph();

        var ptr = Pointer[RefCounted[UnmarshaledType(T)]](pointerState^);
        ptr^.refCount += 1;
        return SharedPointer[UnmarshaledType(T)](ptr), begin;
    }
    case (SHARED_POINTER_DATA) {
        var id = MarshalSize();
        id, begin = ..unmarshalFrom(context, begin, end, MarshalSize);
        if (not null?(lookup(context.pointers, id)))
            throw UnmarshalInconsistentObjectGraph();

        var ptr = allocateRawMemory(RefCounted[UnmarshaledType(T)], MarshalSize(1));
        try {
            ptr^.refCount = 1;
            ptr^.value, begin <-- ..unmarshalFrom(context, begin, end, T);
            context.pointers[id] = RawPointer(ptr);
            return SharedPointer[UnmarshaledType(T)](ptr), begin;
        } catch (ex) {
            freeRawMemory(ptr);
            throw ex;
        }
    }
    else {
        throw UnmarshalInvalidObjectGraphData();
    }
}

[T]
overload UnmarshaledType(static SharedPointer[T]) = SharedPointer[UnmarshaledType(T)];



//
// high-level interface
//
[T | T != MarshalContext] marshal(forward value:T, forward ..values)
    = marshal(MarshalContext(), value, ..values);
overload marshal() = Vector[Byte]();

overload marshal(context: MarshalContext, forward ..values) --> returned: Vector[Byte] {
    var sz = foldValues(add, ..mapValues(v -> marshaledSize(context, v), ..values));
    returned <-- Vector[Byte]();
    resize(returned, sz);

    var buffer = begin(returned);
    ..for (value in values) {
        buffer = marshalTo(context, buffer, value);
    }
    assert(buffer == end(returned));
}

overload marshal(context: MarshalContext) = Vector[Byte]();

unmarshal(data: Vector[Byte], ..Types)
    = ..unmarshal(UnmarshalContext(), data, ..Types);
overload unmarshal(context: UnmarshalContext, data: Vector[Byte], ..Types) {
    var b = begin(data);
    var e = end(data);

    return ..mapValues(T => {
        var n, b2 = ..unmarshalFrom(context, b, e, T);
        b = b2;
        return move(n);
    }, ..Types);
}
