Skip to content

Commit 7b4fd37

Browse files
adam-fowleradam-fowler
authored
Add support for keys stored in argument blocks (#105)
Also simplify constructKeysAffected
1 parent 2354718 commit 7b4fd37

File tree

9 files changed

+60
-55
lines changed

9 files changed

+60
-55
lines changed

‎Sources/Valkey/Commands/BitmapCommands.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -334,7 +334,7 @@ public struct BITOP: ValkeyCommand {
334334
self.key = key
335335
}
336336

337-
public var keysAffected: [ValkeyKey] { [destkey] + key }
337+
public var keysAffected: [ValkeyKey] { key + [destkey] }
338338

339339
@inlinable public func encode(into commandEncoder: inout ValkeyCommandEncoder) {
340340
commandEncoder.encodeArray("BITOP", operation, destkey, key)

‎Sources/Valkey/Commands/GenericCommands.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -710,7 +710,7 @@ public struct SORT: ValkeyCommand {
710710
self.destination = destination
711711
}
712712

713-
public var keysAffected: [ValkeyKey] { [key] + (destination.map { [$0] } ?? []) }
713+
public var keysAffected: [ValkeyKey] { (destination.map { [$0] } ?? []) + [key] }
714714

715715
@inlinable public func encode(into commandEncoder: inout ValkeyCommandEncoder) {
716716
commandEncoder.encodeArray("SORT", key, RESPWithToken("BY", byPattern), RESPWithToken("LIMIT", limit), RESPWithToken("GET", getPattern), order, RESPPureToken("ALPHA", sorting), RESPWithToken("STORE", destination))

‎Sources/Valkey/Commands/HyperloglogCommands.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ public struct PFMERGE: ValkeyCommand {
7070
self.sourcekey = sourcekey
7171
}
7272

73-
public var keysAffected: [ValkeyKey] { [destkey] + sourcekey }
73+
public var keysAffected: [ValkeyKey] { sourcekey + [destkey] }
7474

7575
@inlinable public func encode(into commandEncoder: inout ValkeyCommandEncoder) {
7676
commandEncoder.encodeArray("PFMERGE", destkey, sourcekey)

‎Sources/Valkey/Commands/SetCommands.swift

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ public struct SDIFFSTORE: ValkeyCommand {
9191
self.key = key
9292
}
9393

94-
public var keysAffected: [ValkeyKey] { [destination] + key }
94+
public var keysAffected: [ValkeyKey] { key + [destination] }
9595

9696
@inlinable public func encode(into commandEncoder: inout ValkeyCommandEncoder) {
9797
commandEncoder.encodeArray("SDIFFSTORE", destination, key)
@@ -150,7 +150,7 @@ public struct SINTERSTORE: ValkeyCommand {
150150
self.key = key
151151
}
152152

153-
public var keysAffected: [ValkeyKey] { [destination] + key }
153+
public var keysAffected: [ValkeyKey] { key + [destination] }
154154

155155
@inlinable public func encode(into commandEncoder: inout ValkeyCommandEncoder) {
156156
commandEncoder.encodeArray("SINTERSTORE", destination, key)
@@ -354,7 +354,7 @@ public struct SUNIONSTORE: ValkeyCommand {
354354
self.key = key
355355
}
356356

357-
public var keysAffected: [ValkeyKey] { [destination] + key }
357+
public var keysAffected: [ValkeyKey] { key + [destination] }
358358

359359
@inlinable public func encode(into commandEncoder: inout ValkeyCommandEncoder) {
360360
commandEncoder.encodeArray("SUNIONSTORE", destination, key)

‎Sources/Valkey/Commands/SortedSetCommands.swift

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -250,7 +250,7 @@ public struct ZDIFFSTORE: ValkeyCommand {
250250
self.key = key
251251
}
252252

253-
public var keysAffected: [ValkeyKey] { [destination] + key }
253+
public var keysAffected: [ValkeyKey] { key + [destination] }
254254

255255
@inlinable public func encode(into commandEncoder: inout ValkeyCommandEncoder) {
256256
commandEncoder.encodeArray("ZDIFFSTORE", destination, RESPArrayWithCount(key))
@@ -374,7 +374,7 @@ public struct ZINTERSTORE: ValkeyCommand {
374374
self.aggregate = aggregate
375375
}
376376

377-
public var keysAffected: [ValkeyKey] { [destination] + key }
377+
public var keysAffected: [ValkeyKey] { key + [destination] }
378378

379379
@inlinable public func encode(into commandEncoder: inout ValkeyCommandEncoder) {
380380
commandEncoder.encodeArray("ZINTERSTORE", destination, RESPArrayWithCount(key), RESPWithToken("WEIGHTS", weight), RESPWithToken("AGGREGATE", aggregate))
@@ -1124,7 +1124,7 @@ public struct ZUNIONSTORE: ValkeyCommand {
11241124
self.aggregate = aggregate
11251125
}
11261126

1127-
public var keysAffected: [ValkeyKey] { [destination] + key }
1127+
public var keysAffected: [ValkeyKey] { key + [destination] }
11281128

11291129
@inlinable public func encode(into commandEncoder: inout ValkeyCommandEncoder) {
11301130
commandEncoder.encodeArray("ZUNIONSTORE", destination, RESPArrayWithCount(key), RESPWithToken("WEIGHTS", weight), RESPWithToken("AGGREGATE", aggregate))

‎Sources/Valkey/Commands/StreamCommands.swift

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -627,6 +627,8 @@ public struct XREAD<Id: RESPStringRenderable>: ValkeyCommand {
627627
self.streams = streams
628628
}
629629

630+
public var keysAffected: [ValkeyKey] { streams.key }
631+
630632
public var isBlocking: Bool { true }
631633

632634
public var isReadOnly: Bool { true }
@@ -692,6 +694,8 @@ public struct XREADGROUP<Group: RESPStringRenderable, Consumer: RESPStringRender
692694
self.streams = streams
693695
}
694696

697+
public var keysAffected: [ValkeyKey] { streams.key }
698+
695699
public var isBlocking: Bool { true }
696700

697701
@inlinable public func encode(into commandEncoder: inout ValkeyCommandEncoder) {

‎Sources/Valkey/Commands/StringCommands.swift

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -328,6 +328,8 @@ public struct MSET<Value: RESPStringRenderable>: ValkeyCommand {
328328
self.data = data
329329
}
330330

331+
public var keysAffected: [ValkeyKey] { data.map { $0.key } }
332+
331333
@inlinable public func encode(into commandEncoder: inout ValkeyCommandEncoder) {
332334
commandEncoder.encodeArray("MSET", data)
333335
}
@@ -363,6 +365,8 @@ public struct MSETNX<Value: RESPStringRenderable>: ValkeyCommand {
363365
self.data = data
364366
}
365367

368+
public var keysAffected: [ValkeyKey] { data.map { $0.key } }
369+
366370
@inlinable public func encode(into commandEncoder: inout ValkeyCommandEncoder) {
367371
commandEncoder.encodeArray("MSETNX", data)
368372
}

‎Sources/ValkeyJSON/JsonCommands.swift

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -270,6 +270,8 @@ public enum JSON {
270270
self.data = data
271271
}
272272

273+
public var keysAffected: [ValkeyKey] { data.map { $0.key } }
274+
273275
@inlinable public func encode(into commandEncoder: inout ValkeyCommandEncoder) {
274276
commandEncoder.encodeArray("JSON.MSET", data)
275277
}

‎Sources/_ValkeyCommandsBuilder/ValkeyCommandsRender.swift

Lines changed: 41 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -326,15 +326,14 @@ extension String {
326326
} else {
327327
typeName = name.commandTypeName
328328
}
329-
let keyArguments = command.arguments?.filter { $0.type == .key } ?? []
330329
let conformance = "ValkeyCommand"
331330
let genericTypeParameters = genericTypeParameters(command.arguments)
332331
// Comment header
333332
self.appendCommandCommentHeader(command: command, name: name, tab: tab)
334333
self.append("\(tab)public struct \(typeName)\(genericTypeParameters): \(conformance) {\n")
335334

336335
let arguments = (command.arguments ?? [])
337-
// Enums
336+
// Enums and internal structs
338337
for arg in arguments {
339338
if case .oneOf = arg.type {
340339
self.appendOneOfEnum(argument: arg, names: [], tab: tab)
@@ -374,8 +373,7 @@ extension String {
374373
self.append("\(tab) self.\(arg.name.swiftVariable) = \(arg.name.swiftVariable)\n")
375374
}
376375
self.append("\(tab) }\n\n")
377-
if keyArguments.count > 0 {
378-
let (keysAffectedType, keysAffected) = constructKeysAffected(keyArguments)
376+
if let (keysAffectedType, keysAffected) = constructKeysAffected(command.arguments) {
379377
self.append("\(tab) public var keysAffected: \(keysAffectedType) { \(keysAffected) }\n\n")
380378
}
381379
if command.commandFlags?.contains("BLOCKING") == true {
@@ -492,57 +490,54 @@ private func responseTypeComment(_ response: ValkeyCommand.ReplySchema.Response)
492490
}
493491
}
494492

495-
private func constructKeysAffected(_ keyArguments: [ValkeyCommand.Argument]) -> (type: String, value: String) {
493+
private func constructKeysAffected(_ arguments: [ValkeyCommand.Argument]?) -> (type: String, value: String)? {
494+
guard let arguments else { return nil }
495+
let keyArguments: [(name: String, multiple: Bool)] = arguments.flatMap { getKeyArguments(from: $0) }
496+
guard keyArguments.count > 0 else { return nil }
496497
if keyArguments.count == 1 {
497-
if keyArguments.first!.multiple {
498-
return (type: "[ValkeyKey]", value: keyArguments.first!.name.swiftVariable)
498+
if keyArguments[0].multiple {
499+
return (type: "[ValkeyKey]", value: keyArguments[0].name)
499500
} else {
500-
return (type: "CollectionOfOne<ValkeyKey>", value: ".init(\(keyArguments.first!.name.swiftVariable))")
501+
return (type: "CollectionOfOne<ValkeyKey>", value: ".init(\(keyArguments[0].name))")
501502
}
502503
} else {
503-
var keysAffectedBuilder: String = ""
504-
var inArray = false
505-
var first = true
506-
for key in keyArguments {
507-
if key.multiple {
508-
if inArray {
509-
keysAffectedBuilder += "]"
510-
inArray = false
511-
}
512-
if !first {
513-
keysAffectedBuilder += " + "
514-
}
515-
keysAffectedBuilder += "\(key.name.swiftVariable)"
516-
} else if key.optional {
517-
if inArray {
518-
keysAffectedBuilder += "]"
519-
inArray = false
520-
}
521-
if !first {
522-
keysAffectedBuilder += " + "
523-
}
524-
keysAffectedBuilder += "(\(key.name.swiftVariable).map { [$0] } ?? [])"
525-
} else {
526-
if !inArray {
527-
if !first {
528-
keysAffectedBuilder += " + "
529-
}
530-
keysAffectedBuilder += "[\(key.name.swiftVariable)"
531-
inArray = true
504+
var multipleKeys = keyArguments.compactMap { if $0.multiple { $0.name } else { nil } }
505+
let singleKeys = keyArguments.compactMap { if !$0.multiple { $0.name } else { nil } }
506+
if singleKeys.count > 0 {
507+
multipleKeys.append("[\(singleKeys.joined(separator: ", "))]")
508+
}
509+
let multipleKeysString = multipleKeys.joined(separator: " + ")
510+
return (type: "[ValkeyKey]", value: multipleKeysString)
511+
}
512+
}
513+
514+
private func getKeyArguments(from argument: ValkeyCommand.Argument) -> [(name: String, multiple: Bool)] {
515+
if argument.type == .key {
516+
if argument.multiple {
517+
return [(name: argument.name.swiftVariable, multiple: true)]
518+
} else if argument.optional {
519+
return [(name: "(\(argument.name.swiftVariable).map { [$0] } ?? [])", multiple: true)]
520+
} else {
521+
return [(name: argument.name.swiftVariable, multiple: false)]
522+
}
523+
} else if argument.type == .block, let blockArguments = argument.arguments {
524+
let blockKeyArguments = blockArguments.flatMap { getKeyArguments(from: $0) }
525+
assert(blockKeyArguments.count <= 1, "Do not currently support multiple keys in a block")
526+
if let blockArgument = blockKeyArguments.first {
527+
if argument.multiple {
528+
if blockArgument.multiple {
529+
return [(name: "\(argument.name.swiftVariable).flatMap { $0.\(blockArgument.name) }", multiple: true)]
532530
} else {
533-
if !first {
534-
keysAffectedBuilder += ", "
535-
}
536-
keysAffectedBuilder += "\(key.name.swiftVariable)"
531+
return [(name: "\(argument.name.swiftVariable).map { $0.\(blockArgument.name) }", multiple: true)]
537532
}
533+
} else if argument.optional {
534+
assert(blockKeyArguments.count <= 1, "Do not currently support optional blocks holding keys")
535+
} else {
536+
return [(name: "\(argument.name.swiftVariable).\(blockArgument.name)", multiple: blockArgument.multiple)]
538537
}
539-
first = false
540-
}
541-
if inArray {
542-
keysAffectedBuilder += "]"
543538
}
544-
return (type: "[ValkeyKey]", value: keysAffectedBuilder)
545539
}
540+
return []
546541
}
547542

548543
private func subCommand(_ command: String) -> (String.SubSequence, String.SubSequence?) {

0 commit comments

Comments
 (0)