a size primitive for sets

Reviewed By: donsbot Differential Revision: D67405373 fbshipit-source-id: dc0167fb8571a8325df34ffa53cc23d08a591d4d
facebookincubator · Dec 19, 2024 · 6b87d5f · 6b87d5f
1 parent 9e6e4b9
commit 6b87d5f
Show file tree

Hide file tree

Showing 6 changed files with 22 additions and 1 deletion.
diff --git a/glean/db/Glean/Query/Codegen.hs b/glean/db/Glean/Query/Codegen.hs
@@ -742,6 +742,12 @@ compileStatements
             getOutput array ptr end
             inputNat ptr end reg
           inner
+      compileGen (PrimCall PrimOpSize [arg] _) (Just reg) inner =
+        withTerm vars arg $ \set -> do
+          local $ \ptr end -> do
+            getOutput set ptr end
+            inputNat ptr end reg
+          inner
       compileGen (PrimCall PrimOpZip [arg1,arg2]
                     (Angle.ArrayTy (Angle.RecordTy fields)))
                   (Just reg) inner

diff --git a/glean/db/Glean/Query/Typecheck.hs b/glean/db/Glean/Query/Typecheck.hs
@@ -928,6 +928,9 @@ primOpType op = case op of
     x <- freshTyVar
     return ([ArrayTy x, ArrayTy x], ArrayTy x)
   PrimOpReverse -> return  ([StringTy], StringTy)
+  PrimOpSize -> do
+    x <- freshTyVar
+    return ([SetTy x], NatTy)
   PrimOpRelToAbsByteSpans ->
     -- prim.relToAbsByteSpans takes an array of pairs as input and
     -- returns an array of pairs as output

diff --git a/glean/hs/Glean/Angle/Types.hs b/glean/hs/Glean/Angle/Types.hs
@@ -338,6 +338,7 @@ data PrimOp
   | PrimOpZip
   | PrimOpConcat
   | PrimOpReverse
+  | PrimOpSize
   | PrimOpRelToAbsByteSpans
   | PrimOpUnpackByteSpans
   | PrimOpGtNat
@@ -676,12 +677,14 @@ newtype AngleVersion = AngleVersion Int
 --   - added dot syntax
 -- version 10:
 --   - optional versions in `import`, `evolves`, inheritance
+-- version 11:
+--   - size primop
 
 latestSupportedAngleVersion :: AngleVersion
 latestSupportedAngleVersion = AngleVersion 5
 
 latestAngleVersion :: AngleVersion
-latestAngleVersion = AngleVersion 10
+latestAngleVersion = AngleVersion 11
 
 -- -----------------------------------------------------------------------------
 -- Pretty-printing
@@ -971,6 +974,7 @@ instance Display PrimOp where
   display _ PrimOpZip = "prim.zip"
   display _ PrimOpConcat = "prim.concat"
   display _ PrimOpReverse = "prim.reverse"
+  display _ PrimOpSize = "prim.size"
   display _ PrimOpRelToAbsByteSpans = "prim.relToAbsByteSpans"
   display _ PrimOpUnpackByteSpans = "prim.unpackByteSpans"
   display _ PrimOpGtNat = "prim.gtNat"

diff --git a/glean/hs/Glean/Schema/Resolve.hs b/glean/hs/Glean/Schema/Resolve.hs
@@ -801,6 +801,7 @@ primitives = HashMap.fromList
   , ("prim.zip", PrimOpZip)
   , ("prim.concat", PrimOpConcat)
   , ("prim.reverse", PrimOpReverse)
+  , ("prim.size", PrimOpSize)
   , ("prim.gtNat", PrimOpGtNat)
   , ("prim.geNat", PrimOpGeNat)
   , ("prim.ltNat", PrimOpLtNat)

diff --git a/glean/test/tests/Angle/SetTest.hs b/glean/test/tests/Angle/SetTest.hs
@@ -121,3 +121,6 @@ setSemanticsTest = dbTestCase $ \env repo -> do
 
   r <- runQuery_ env repo $ angleData @Nat [s| elements (all 65535) |]
   assertEqual "big nat" r [Nat 65535]
+
+  r <- runQuery_ env repo $ angleData @Nat [s| prim.size (all (1 | 2 | 3))|]
+  assertEqual "size" r [Nat 3]
diff --git a/glean/website/docs/angle/reference.md b/glean/website/docs/angle/reference.md
@@ -169,6 +169,10 @@ Angle supports a few primitive operations. The argument(s) to a primitive operat
 
 > Equal to the number of elements in its array argument
 
+&nbsp;&nbsp;`prim.size` (A : set _) : nat
+
+> Equal to the number of elements in the set
+
 &nbsp;&nbsp;*term* `>` *term* <br />
 &nbsp;&nbsp;*term* `>=` *term* <br />
 &nbsp;&nbsp;*term* `<` *term* <br />