import Foundation private enum MediaEntryType: Int8 { case Direct case MessageReference } enum InsertMediaResult { case Reference case Embed(Media) } enum RemoveMediaResult { case Reference case Embedded(MessageIndex) } enum DebugMediaEntry { case Direct(Media, Int) case MessageReference(MessageIndex) } final class MessageMediaTable: Table { static func tableSpec(_ id: Int32) -> ValueBoxTable { return ValueBoxTable(id: id, keyType: .binary, compactValuesOnCreation: false) } func key(_ id: MediaId, key: ValueBoxKey = ValueBoxKey(length: 4 + 8)) -> ValueBoxKey { key.setInt32(0, value: id.namespace) key.setInt64(4, value: id.id) return key } func get(_ id: MediaId, embedded: (MessageIndex, MediaId) -> Media?) -> (MessageIndex?, Media)? { if let value = self.valueBox.get(self.table, key: self.key(id)) { var type: Int8 = 0 value.read(&type, offset: 0, length: 1) if type == MediaEntryType.Direct.rawValue { var dataLength: Int32 = 0 value.read(&dataLength, offset: 0, length: 4) if let media = PostboxDecoder(buffer: MemoryBuffer(memory: value.memory + value.offset, capacity: Int(dataLength), length: Int(dataLength), freeWhenDone: false)).decodeRootObject() as? Media { return (nil, media) } } else if type == MediaEntryType.MessageReference.rawValue { var idPeerId: Int64 = 0 var idNamespace: Int32 = 0 var idId: Int32 = 0 var idTimestamp: Int32 = 0 value.read(&idPeerId, offset: 0, length: 8) value.read(&idNamespace, offset: 0, length: 4) value.read(&idId, offset: 0, length: 4) value.read(&idTimestamp, offset: 0, length: 4) let referencedMessageIndex = MessageIndex(id: MessageId(peerId: PeerId(idPeerId), namespace: idNamespace, id: idId), timestamp: idTimestamp) if let result = embedded(referencedMessageIndex, id) { return (referencedMessageIndex, result) } else { return nil } } } return nil } func set(_ media: Media, index: MessageIndex, messageHistoryTable: MessageHistoryTable, sharedWriteBuffer: WriteBuffer = WriteBuffer(), sharedEncoder: PostboxEncoder = PostboxEncoder()) -> InsertMediaResult { if let id = media.id { if let value = self.valueBox.get(self.table, key: self.key(id)) { var type: Int8 = 0 value.read(&type, offset: 0, length: 1) if type == MediaEntryType.Direct.rawValue { var dataLength: Int32 = 0 value.read(&dataLength, offset: 0, length: 4) value.skip(Int(dataLength)) sharedWriteBuffer.reset() sharedWriteBuffer.write(value.memory, offset: 0, length: value.offset) var messageReferenceCount: Int32 = 0 value.read(&messageReferenceCount, offset: 0, length: 4) messageReferenceCount += 1 sharedWriteBuffer.write(&messageReferenceCount, offset: 0, length: 4) withExtendedLifetime(sharedWriteBuffer, { self.valueBox.set(self.table, key: self.key(id), value: sharedWriteBuffer.readBufferNoCopy()) }) return .Reference } else if type == MediaEntryType.MessageReference.rawValue { var idPeerId: Int64 = 0 var idNamespace: Int32 = 0 var idId: Int32 = 0 var idTimestamp: Int32 = 0 value.read(&idPeerId, offset: 0, length: 8) value.read(&idNamespace, offset: 0, length: 4) value.read(&idId, offset: 0, length: 4) value.read(&idTimestamp, offset: 0, length: 4) let referencedMessageIndex = MessageIndex(id: MessageId(peerId: PeerId(idPeerId), namespace: idNamespace, id: idId), timestamp: idTimestamp) if referencedMessageIndex == index { return .Embed(media) } if let media = messageHistoryTable.unembedMedia(referencedMessageIndex, id: id) { sharedWriteBuffer.reset() var directType: Int8 = MediaEntryType.Direct.rawValue sharedWriteBuffer.write(&directType, offset: 0, length: 1) sharedEncoder.reset() sharedEncoder.encodeRootObject(media) let mediaBuffer = sharedEncoder.memoryBuffer() var mediaBufferLength = Int32(mediaBuffer.length) sharedWriteBuffer.write(&mediaBufferLength, offset: 0, length: 4) sharedWriteBuffer.write(mediaBuffer.memory, offset: 0, length: mediaBuffer.length) var messageReferenceCount: Int32 = 2 sharedWriteBuffer.write(&messageReferenceCount, offset: 0, length: 4) withExtendedLifetime(sharedWriteBuffer, { self.valueBox.set(self.table, key: self.key(id), value: sharedWriteBuffer.readBufferNoCopy()) }) } return .Reference } else { return .Embed(media) } } else { sharedWriteBuffer.reset() var type: Int8 = MediaEntryType.MessageReference.rawValue sharedWriteBuffer.write(&type, offset: 0, length: 1) var idPeerId: Int64 = index.id.peerId.toInt64() var idNamespace: Int32 = index.id.namespace var idId: Int32 = index.id.id var idTimestamp: Int32 = index.timestamp sharedWriteBuffer.write(&idPeerId, offset: 0, length: 8) sharedWriteBuffer.write(&idNamespace, offset: 0, length: 4) sharedWriteBuffer.write(&idId, offset: 0, length: 4) sharedWriteBuffer.write(&idTimestamp, offset: 0, length: 4) withExtendedLifetime(sharedWriteBuffer, { self.valueBox.set(self.table, key: self.key(id), value: sharedWriteBuffer.readBufferNoCopy()) }) return .Embed(media) } } else { return .Embed(media) } } func removeReference(_ id: MediaId, sharedWriteBuffer: WriteBuffer = WriteBuffer()) -> RemoveMediaResult { if let value = self.valueBox.get(self.table, key: self.key(id)) { var type: Int8 = 0 value.read(&type, offset: 0, length: 1) if type == MediaEntryType.Direct.rawValue { var dataLength: Int32 = 0 value.read(&dataLength, offset: 0, length: 4) value.skip(Int(dataLength)) sharedWriteBuffer.reset() sharedWriteBuffer.write(value.memory, offset: 0, length: value.offset) var messageReferenceCount: Int32 = 0 value.read(&messageReferenceCount, offset: 0, length: 4) messageReferenceCount -= 1 sharedWriteBuffer.write(&messageReferenceCount, offset: 0, length: 4) if messageReferenceCount <= 0 { self.valueBox.remove(self.table, key: self.key(id), secure: false) } else { withExtendedLifetime(sharedWriteBuffer, { self.valueBox.set(self.table, key: self.key(id), value: sharedWriteBuffer.readBufferNoCopy()) }) } return .Reference } else if type == MediaEntryType.MessageReference.rawValue { var idPeerId: Int64 = 0 var idNamespace: Int32 = 0 var idId: Int32 = 0 var idTimestamp: Int32 = 0 value.read(&idPeerId, offset: 0, length: 8) value.read(&idNamespace, offset: 0, length: 4) value.read(&idId, offset: 0, length: 4) value.read(&idTimestamp, offset: 0, length: 4) let referencedMessageIndex = MessageIndex(id: MessageId(peerId: PeerId(idPeerId), namespace: idNamespace, id: idId), timestamp: idTimestamp) self.valueBox.remove(self.table, key: self.key(id), secure: false) return .Embedded(referencedMessageIndex) } else { assertionFailure() } } return .Reference } func removeEmbeddedMedia(_ media: Media) { if let id = media.id { self.valueBox.remove(self.table, key: self.key(id), secure: false) } } func update(_ id: MediaId, media: Media, messageHistoryTable: MessageHistoryTable, operationsByPeerId: inout [PeerId: [MessageHistoryOperation]], sharedWriteBuffer: WriteBuffer = WriteBuffer(), sharedEncoder: PostboxEncoder = PostboxEncoder()) { if let updatedId = media.id { if let value = self.valueBox.get(self.table, key: self.key(id)) { var type: Int8 = 0 value.read(&type, offset: 0, length: 1) if type == MediaEntryType.Direct.rawValue { var dataLength: Int32 = 0 value.read(&dataLength, offset: 0, length: 4) value.skip(Int(dataLength)) var messageReferenceCount: Int32 = 0 value.read(&messageReferenceCount, offset: 0, length: 4) sharedWriteBuffer.reset() var directType: Int8 = MediaEntryType.Direct.rawValue sharedWriteBuffer.write(&directType, offset: 0, length: 1) sharedEncoder.reset() sharedEncoder.encodeRootObject(media) let mediaBuffer = sharedEncoder.memoryBuffer() var mediaBufferLength = Int32(mediaBuffer.length) sharedWriteBuffer.write(&mediaBufferLength, offset: 0, length: 4) sharedWriteBuffer.write(mediaBuffer.memory, offset: 0, length: mediaBuffer.length) sharedWriteBuffer.write(&messageReferenceCount, offset: 0, length: 4) if id != updatedId { self.valueBox.remove(self.table, key: self.key(id), secure: false) } withExtendedLifetime(sharedWriteBuffer, { self.valueBox.set(self.table, key: self.key(updatedId), value: sharedWriteBuffer.readBufferNoCopy()) }) } else if type == MediaEntryType.MessageReference.rawValue { var idPeerId: Int64 = 0 var idNamespace: Int32 = 0 var idId: Int32 = 0 var idTimestamp: Int32 = 0 value.read(&idPeerId, offset: 0, length: 8) value.read(&idNamespace, offset: 0, length: 4) value.read(&idId, offset: 0, length: 4) value.read(&idTimestamp, offset: 0, length: 4) let referencedMessageIndex = MessageIndex(id: MessageId(peerId: PeerId(idPeerId), namespace: idNamespace, id: idId), timestamp: idTimestamp) messageHistoryTable.updateEmbeddedMedia(referencedMessageIndex, mediaId: id, media: media, operationsByPeerId: &operationsByPeerId) } } } } func debugList() -> [DebugMediaEntry] { var entries: [DebugMediaEntry] = [] let upperBoundKey = ValueBoxKey(length: 8 + 4) memset(upperBoundKey.memory, 0xff, 8 + 4) self.valueBox.range(self.table, start: ValueBoxKey(length: 0), end: upperBoundKey, values: { key, value in var type: Int8 = 0 value.read(&type, offset: 0, length: 1) if type == MediaEntryType.Direct.rawValue { var dataLength: Int32 = 0 value.read(&dataLength, offset: 0, length: 4) if let media = PostboxDecoder(buffer: MemoryBuffer(memory: value.memory + value.offset, capacity: Int(dataLength), length: Int(dataLength), freeWhenDone: false)).decodeRootObject() as? Media { value.skip(Int(dataLength)) var messageReferenceCount: Int32 = 0 value.read(&messageReferenceCount, offset: 0, length: 4) entries.append(.Direct(media, Int(messageReferenceCount))) } } else if type == MediaEntryType.MessageReference.rawValue { var idPeerId: Int64 = 0 var idNamespace: Int32 = 0 var idId: Int32 = 0 var idTimestamp: Int32 = 0 value.read(&idPeerId, offset: 0, length: 8) value.read(&idNamespace, offset: 0, length: 4) value.read(&idId, offset: 0, length: 4) value.read(&idTimestamp, offset: 0, length: 4) let referencedMessageIndex = MessageIndex(id: MessageId(peerId: PeerId(idPeerId), namespace: idNamespace, id: idId), timestamp: idTimestamp) entries.append(.MessageReference(referencedMessageIndex)) } return true }, limit: 1000) return entries } }