Skip to content

Commit b277bb6

Browse files
lihaoyi-databrickslihaoyi-databricks
authored
Merge pull request #40 from modusintegration/Key-Functions
Key functions
2 parents 7a184dc + 4bc6745 commit b277bb6

3 files changed

Lines changed: 539 additions & 106 deletions

File tree

‎build.sc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -136,4 +136,4 @@ class SjsonnetModule(val crossScalaVersion: String) extends Module {
136136
}
137137
}
138138

139-
}
139+
}

‎sjsonnet/src/sjsonnet/Std.scala‎

Lines changed: 216 additions & 105 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,14 @@ import java.io.StringWriter
44
import java.util.Base64
55

66
import sjsonnet.Expr.Member.Visibility
7-
import sjsonnet.Expr.Params
7+
import sjsonnet.Expr.{BinaryOp, False, Params}
88

99
import scala.collection.mutable.ArrayBuffer
1010
import scala.collection.compat._
11+
import sjsonnet.Std.builtinWithDefaults
12+
import ujson.Value
13+
14+
import util.control.Breaks._
1115

1216
/**
1317
* The Jsonnet standard library, `std`, with each builtin function implemented
@@ -478,135 +482,159 @@ object Std {
478482
builtin("base64DecodeBytes", "s"){ (ev, fs, s: String) =>
479483
Val.Arr(Base64.getDecoder().decode(s).map(i => Val.Lazy(Val.Num(i))))
480484
},
481-
builtin("sort", "arr"){ (ev, fs, arr: Val) =>
482-
arr match{
483-
case Val.Arr(vs) =>
484-
Val.Arr(
485-
486-
if (vs.forall(_.force.isInstanceOf[Val.Str])){
487-
vs.map(_.force.cast[Val.Str]).sortBy(_.value).map(Val.Lazy(_))
488-
}else if (vs.forall(_.force.isInstanceOf[Val.Num])){
489-
vs.map(_.force.cast[Val.Num]).sortBy(_.value).map(Val.Lazy(_))
490-
}else {
491-
???
492-
}
493-
)
494-
case Val.Str(s) => Val.Arr(s.sorted.map(c => Val.Lazy(Val.Str(c.toString))))
495-
case x => throw new Error.Delegate("Cannot sort " + x.prettyName)
496-
}
497-
},
498-
builtin("uniq", "arr"){ (ev, fs, arr: Val.Arr) =>
499-
val ujson.Arr(vs) = Materializer(arr)(ev)
500-
val out = collection.mutable.Buffer.empty[ujson.Value]
501-
for(v <- vs) if (out.isEmpty || out.last != v) out.append(v)
502-
503-
Val.Arr(out.map(v => Val.Lazy(Materializer.reverse(v))).toSeq)
504-
},
505-
builtin("set", "arr"){ (ev, fs, arr: Val.Arr) =>
506-
val ujson.Arr(vs0) = Materializer(arr)(ev)
507-
val vs =
508-
if (vs0.forall(_.isInstanceOf[ujson.Str])){
509-
vs0.map(_.asInstanceOf[ujson.Str]).sortBy(_.value)
510-
}else if (vs0.forall(_.isInstanceOf[ujson.Num])){
511-
vs0.map(_.asInstanceOf[ujson.Num]).sortBy(_.value)
512-
}else {
513-
throw new Error.Delegate("Every element of the input must be of the same type, string or number")
514-
}
515485

516-
val out = collection.mutable.Buffer.empty[ujson.Value]
517-
for(v <- vs) if (out.isEmpty || out.last != v) out.append(v)
486+
builtinWithDefaults("uniq", "arr" -> None, "keyF" -> Some(Expr.False(0))) { (args, ev) =>
487+
val arr = args("arr")
488+
val keyF = args("keyF")
518489

519-
Val.Arr(out.map(v => Val.Lazy(Materializer.reverse(v))).toSeq)
490+
uniqArr(ev, arr, keyF)
520491
},
521-
builtin("setUnion", "a", "b"){ (ev, fs, a: Val.Arr, b: Val.Arr) =>
492+
builtinWithDefaults("sort", "arr" -> None, "keyF" -> Some(Expr.False(0))) { (args, ev) =>
493+
val arr = args("arr")
494+
val keyF = args("keyF")
522495

523-
val ujson.Arr(vs1) = Materializer(a)(ev)
524-
val ujson.Arr(vs2) = Materializer(b)(ev)
525-
val vs0 = vs1 ++ vs2
526-
val vs =
527-
if (vs0.forall(_.isInstanceOf[ujson.Str])){
528-
vs0.map(_.asInstanceOf[ujson.Str]).sortBy(_.value)
529-
}else if (vs0.forall(_.isInstanceOf[ujson.Num])){
530-
vs0.map(_.asInstanceOf[ujson.Num]).sortBy(_.value)
531-
}else {
532-
throw new Error.Delegate("Every element of the input must be of the same type, string or number")
533-
}
496+
sortArr(ev, arr, keyF)
497+
},
534498

535-
val out = collection.mutable.Buffer.empty[ujson.Value]
536-
for(v <- vs) if (out.isEmpty || out.last != v) out.append(v)
499+
builtinWithDefaults("set", "arr" -> None, "keyF" -> Some(Expr.False(0))) { (args, ev) =>
500+
uniqArr(ev, sortArr(ev, args("arr"), args("keyF")), args("keyF"))
501+
},
502+
builtinWithDefaults("setUnion", "a" -> None, "b" -> None, "keyF" -> Some(Expr.False(0))) { (args, ev) =>
503+
val a = args("a") match {
504+
case arr: Val.Arr => arr.value
505+
case str: Val.Str => stringChars(str.value).value
506+
case _ => throw new Error.Delegate("Arguments must be either arrays or strings")
507+
}
508+
val b = args("b") match {
509+
case arr: Val.Arr => arr.value
510+
case str: Val.Str => stringChars(str.value).value
511+
case _ => throw new Error.Delegate("Arguments must be either arrays or strings")
512+
}
537513

538-
Val.Arr(out.map(v => Val.Lazy(Materializer.reverse(v))).toSeq)
514+
val concat = Val.Arr(a ++ b)
515+
uniqArr(ev, sortArr(ev, concat, args("keyF")), args("keyF"))
539516
},
540-
builtin("setInter", "a", "b"){ (ev, fs, a: Val, b: Val.Arr) =>
541-
val vs1 = Materializer(a)(ev) match{
542-
case ujson.Arr(vs1) => vs1
543-
case x => Seq(x)
517+
builtinWithDefaults("setInter", "a" -> None, "b" -> None, "keyF" -> Some(Expr.False(0))) { (args, ev) =>
518+
val a = args("a") match {
519+
case arr: Val.Arr => arr.value
520+
case str: Val.Str => stringChars(str.value).value
521+
case _ => throw new Error.Delegate("Arguments must be either arrays or strings")
522+
}
523+
val b = args("b") match {
524+
case arr: Val.Arr => arr.value
525+
case str: Val.Str => stringChars(str.value).value
526+
case _ => throw new Error.Delegate("Arguments must be either arrays or strings")
544527
}
545-
val ujson.Arr(vs2) = Materializer(b)(ev)
546528

529+
val keyF = args("keyF")
530+
val out = collection.mutable.Buffer.empty[Val.Lazy]
547531

548-
val vs0 = vs1.to(collection.mutable.LinkedHashSet)
549-
.intersect(vs2.to(collection.mutable.LinkedHashSet))
550-
.toSeq
551-
val vs =
552-
if (vs0.forall(_.isInstanceOf[ujson.Str])){
553-
vs0.map(_.asInstanceOf[ujson.Str]).sortBy(_.value)
554-
}else if (vs0.forall(_.isInstanceOf[ujson.Num])){
555-
vs0.map(_.asInstanceOf[ujson.Num]).sortBy(_.value)
556-
}else {
557-
throw new Error.Delegate("Every element of the input must be of the same type, string or number")
532+
for (v <- a) {
533+
if (keyF == Val.False) {
534+
val mv = Materializer.apply(v.force)(ev)
535+
if (b.exists(value => {
536+
val mValue = Materializer.apply(value.force)(ev)
537+
mValue == mv
538+
}) && !out.exists(value => {
539+
val mValue = Materializer.apply(value.force)(ev)
540+
mValue == mv
541+
})) {
542+
out.append(v)
543+
}
544+
} else {
545+
val keyFFunc = keyF.asInstanceOf[Val.Func]
546+
val keyFApplyer = Applyer(keyFFunc, ev, null)
547+
val appliedX = keyFApplyer.apply(v)
548+
549+
if (b.exists(value => {
550+
val appliedValue = keyFApplyer.apply(value)
551+
appliedValue == appliedX
552+
}) && !out.exists(value => {
553+
val mValue = keyFApplyer.apply(value)
554+
mValue == appliedX
555+
})) {
556+
out.append(v)
557+
}
558558
}
559+
}
559560

560-
val out = collection.mutable.Buffer.empty[ujson.Value]
561-
for(v <- vs) if (out.isEmpty || out.last != v) out.append(v)
562-
563-
Val.Arr(out.map(v => Val.Lazy(Materializer.reverse(v))).toSeq)
561+
sortArr(ev, Val.Arr(out.toSeq), keyF)
564562
},
565-
builtin("setDiff", "a", "b"){ (ev, fs, a: Val.Arr, b: Val.Arr) =>
566-
val ujson.Arr(vs1) = Materializer(a)(ev)
567-
val ujson.Arr(vs2) = Materializer(b)(ev)
563+
builtinWithDefaults("setDiff", "a" -> None, "b" -> None, "keyF" -> Some(Expr.False(0))) { (args, ev) =>
568564

565+
val a = args("a") match {
566+
case arr: Val.Arr => arr.value
567+
case str: Val.Str => stringChars(str.value).value
568+
case _ => throw new Error.Delegate("Arguments must be either arrays or strings")
569+
}
570+
val b = args("b") match {
571+
case arr: Val.Arr => arr.value
572+
case str: Val.Str => stringChars(str.value).value
573+
case _ => throw new Error.Delegate("Arguments must be either arrays or strings")
574+
}
569575

570-
val vs0 = vs1.to(collection.mutable.LinkedHashSet)
571-
.diff(vs2.to(collection.mutable.LinkedHashSet))
572-
.toSeq
573-
val vs =
574-
if (vs0.forall(_.isInstanceOf[ujson.Str])){
575-
vs0.map(_.asInstanceOf[ujson.Str]).sortBy(_.value)
576-
}else if (vs0.forall(_.isInstanceOf[ujson.Num])){
577-
vs0.map(_.asInstanceOf[ujson.Num]).sortBy(_.value)
578-
}else {
579-
throw new Error.Delegate("Every element of the input must be of the same type, string or number")
580-
}
581-
582-
val out = collection.mutable.Buffer.empty[ujson.Value]
583-
for(v <- vs) if (out.isEmpty || out.last != v) out.append(v)
576+
val keyF = args("keyF")
577+
val out = collection.mutable.Buffer.empty[Val.Lazy]
584578

585-
Val.Arr(out.map(v => Val.Lazy(Materializer.reverse(v))).toSeq)
579+
for (v <- a) {
580+
if (keyF == Val.False) {
581+
val mv = Materializer.apply(v.force)(ev)
582+
if (!b.exists(value => {
583+
val mValue = Materializer.apply(value.force)(ev)
584+
mValue == mv
585+
}) && !out.exists(value => {
586+
val mValue = Materializer.apply(value.force)(ev)
587+
mValue == mv
588+
})) {
589+
out.append(v)
590+
}
591+
} else {
592+
val keyFFunc = keyF.asInstanceOf[Val.Func]
593+
val keyFApplyer = Applyer(keyFFunc, ev, null)
594+
val appliedX = keyFApplyer.apply(v)
595+
596+
if (!b.exists(value => {
597+
val appliedValue = keyFApplyer.apply(value)
598+
appliedValue == appliedX
599+
}) && !out.exists(value => {
600+
val mValue = keyFApplyer.apply(value)
601+
mValue == appliedX
602+
})) {
603+
out.append(v)
604+
}
605+
}
606+
}
586607

608+
sortArr(ev, Val.Arr(out.toSeq), keyF)
609+
},
610+
builtinWithDefaults("setMember", "x" -> None, "arr" -> None, "keyF" -> Some(Expr.False(0))) { (args, ev) =>
611+
val keyF = args("keyF")
612+
613+
if (keyF == Val.False) {
614+
val ujson.Arr(mArr) = Materializer(args("arr"))(ev)
615+
val mx = Materializer(args("x"))(ev)
616+
mArr.contains(mx)
617+
} else {
618+
val x = Val.Lazy(args("x"))
619+
val arr = args("arr").asInstanceOf[Val.Arr].value
620+
val keyFFunc = keyF.asInstanceOf[Val.Func]
621+
val keyFApplyer = Applyer(keyFFunc, ev, null)
622+
val appliedX = keyFApplyer.apply(x)
623+
arr.exists(value => {
624+
val appliedValue = keyFApplyer.apply(value)
625+
appliedValue == appliedX
626+
})
627+
}
587628
},
588-
builtin("setMember", "x", "arr"){ (ev, fs, x: Val, arr: Val.Arr) =>
589-
val vs1 = Materializer(x)(ev)
590-
val ujson.Arr(vs2) = Materializer(arr)(ev)
591-
vs2.contains(vs1)
592-
},
629+
593630
builtin("split", "str", "c"){ (ev, fs, str: String, c: String) =>
594631
Val.Arr(str.split(java.util.regex.Pattern.quote(c), -1).map(s => Val.Lazy(Val.Str(s))))
595632
},
596633
builtin("splitLimit", "str", "c", "maxSplits"){ (ev, fs, str: String, c: String, maxSplits: Int) =>
597634
Val.Arr(str.split(java.util.regex.Pattern.quote(c), maxSplits + 1).map(s => Val.Lazy(Val.Str(s))))
598635
},
599636
builtin("stringChars", "str"){ (ev, fs, str: String) =>
600-
601-
var offset = 0
602-
val output = collection.mutable.Buffer.empty[String]
603-
while (offset < str.length) {
604-
val codepoint = str.codePointAt(offset)
605-
output.append(new String(Character.toChars(codepoint)))
606-
offset += Character.charCount(codepoint)
607-
}
608-
Val.Arr(output.map(s => Val.Lazy(Val.Str(s))).toSeq)
609-
637+
stringChars(str)
610638
},
611639
builtin("parseInt", "str"){ (ev, fs, str: String) =>
612640
str.toInt
@@ -676,6 +704,7 @@ object Std {
676704
scope.bindings(1).get.force
677705
}
678706
),
707+
679708
"extVar" -> Val.Func(
680709
None,
681710
Params(Array(("x", None, 0))),
@@ -791,4 +820,86 @@ object Std {
791820
None, None, None, Array(Val.Lazy(Std)).padTo(size, null)
792821
)
793822
}
823+
824+
def uniqArr(ev: EvalScope, arr: Val, keyF: Val) = {
825+
val arrValue = arr match {
826+
case arr: Val.Arr => arr.value
827+
case str: Val.Str => stringChars(str.value).value
828+
case _ => throw new Error.Delegate("Argument must be either array or string")
829+
}
830+
831+
val out = collection.mutable.Buffer.empty[Val.Lazy]
832+
for (v <- arrValue) {
833+
if (out.isEmpty) {
834+
out.append(v)
835+
} else if (keyF == Val.False) {
836+
val ol = Materializer.apply(out.last.force)(ev)
837+
val mv = Materializer.apply(v.force)(ev)
838+
if (ol != mv) {
839+
out.append(v)
840+
}
841+
} else if (keyF != Val.False) {
842+
val keyFFunc = keyF.asInstanceOf[Val.Func]
843+
val keyFApplyer = Applyer(keyFFunc, ev, null)
844+
845+
val o1Key = keyFApplyer.apply(v)
846+
val o2Key = keyFApplyer.apply(out.last)
847+
val o1KeyExpr = Materializer.toExpr(Materializer.apply(o1Key)(ev))
848+
val o2KeyExpr = Materializer.toExpr(Materializer.apply(o2Key)(ev))
849+
850+
val comparisonExpr = Expr.BinaryOp(0, o1KeyExpr, BinaryOp.`!=`, o2KeyExpr)
851+
val exprResult = ev.visitExpr(comparisonExpr)(scope(0), new FileScope(null, Map.empty))
852+
853+
val res = Materializer.apply(exprResult)(ev).asInstanceOf[ujson.Bool]
854+
855+
if (res.value) {
856+
out.append(v)
857+
}
858+
}
859+
}
860+
861+
Val.Arr(out.toSeq)
862+
}
863+
864+
def sortArr(ev: EvalScope, arr: Val, keyF: Val) = {
865+
arr match{
866+
case Val.Arr(vs) =>
867+
Val.Arr(
868+
869+
if (vs.forall(_.force.isInstanceOf[Val.Str])){
870+
vs.map(_.force.cast[Val.Str]).sortBy(_.value).map(Val.Lazy(_))
871+
}else if (vs.forall(_.force.isInstanceOf[Val.Num])) {
872+
vs.map(_.force.cast[Val.Num]).sortBy(_.value).map(Val.Lazy(_))
873+
}else if (vs.forall(_.force.isInstanceOf[Val.Obj])){
874+
if (keyF == Val.False) {
875+
throw new Error.Delegate("Unable to sort array of objects without key function")
876+
} else {
877+
val keyFFunc = keyF.asInstanceOf[Val.Func]
878+
val keyFApplyer = Applyer(keyFFunc, ev, null)
879+
vs.map(_.force.cast[Val.Obj]).sortWith((o1, o2) => {
880+
val o1Key = keyFApplyer.apply(Val.Lazy(o1))
881+
val o2Key = keyFApplyer.apply(Val.Lazy(o2))
882+
val o1KeyExpr = Materializer.toExpr(Materializer.apply(o1Key)(ev))
883+
val o2KeyExpr = Materializer.toExpr(Materializer.apply(o2Key)(ev))
884+
885+
val comparisonExpr = Expr.BinaryOp(0, o1KeyExpr, BinaryOp.`<`, o2KeyExpr)
886+
val exprResult = ev.visitExpr(comparisonExpr)(scope(0), new FileScope(null, Map.empty))
887+
val res = Materializer.apply(exprResult)(ev).asInstanceOf[ujson.Bool]
888+
res.value
889+
}).map(Val.Lazy(_))
890+
}
891+
}else {
892+
???
893+
}
894+
)
895+
case Val.Str(s) => Val.Arr(s.sorted.map(c => Val.Lazy(Val.Str(c.toString))))
896+
case x => throw new Error.Delegate("Cannot sort " + x.prettyName)
897+
}
898+
}
899+
900+
def stringChars(str: String): Val.Arr = {
901+
var offset = 0
902+
val output = str.toSeq.sliding(1).toList
903+
Val.Arr(output.map(s => Val.Lazy(Val.Str(s.toString()))).toSeq)
904+
}
794905
}

0 commit comments

Comments
 (0)