Ancient Haskell version of jq. Might be useful someday. Maybe.haskell-version

7 files changed, 269 insertions, 46 deletions

diff --git a/JQ.hs b/JQ.hs
index ca8df794..a8a1e4f3 100644
--- a/JQ.hs
+++ b/JQ.hs
@@ -2,55 +2,93 @@ module JQ where
 import Text.JSON
 import Text.JSON.String
 import Data.Maybe
-import Data.List (sortBy,sort,groupBy)
+import Data.Char
+import Data.List (sortBy,sort,groupBy,partition,intercalate)
 import Data.Function (on)
 import Data.Ord (comparing)
 import Control.Monad
 import Control.Monad.Writer
 import Control.Monad.List
 import Control.Monad.Reader
+import qualified Data.Map as M
+import qualified Data.Set as S
+import Data.Map ((!))
+import Debug.Trace
 
 type Path = [Either Int String]
 
 type Program = JSValue -> [(JSValue, Path)]
 
-type JQ = ReaderT JSValue (WriterT Path [])
+type Filter = JQ JSValue
+newtype Operator = Operator {runOperator:: [Filter] -> Filter}
+
+type JQ = ReaderT (JSValue, M.Map String JSValue, M.Map (String, Int) Operator) (WriterT Path [])
 
 runJQ :: JQ a -> JSValue -> [a]
-runJQ prog val = map fst $ runWriterT $ runReaderT prog val
+runJQ prog val = map fst $ runWriterT $ runReaderT prog (val,M.empty,primitives)
 
 (>|) :: JQ JSValue -> JQ a -> JQ a
 a >| b = do
   val <- a
-  local (const val) b
+  local (\(v,s,d) -> (val,s,d)) b
+
+setvar name val prog = 
+  local (\(v,s,d) -> (v, M.insert name val s, d)) prog
+getvar name = liftM (! name) $ asks (\(v,s,d) -> s)
+
+input = asks (\(v,s,d) -> v)
 
 collect :: JQ a -> JQ [a]
-collect prog = do
-  arg <- ask
-  return $ runJQ prog arg
+collect prog = liftM (map fst) $ collectPaths prog
   
 collectPaths :: JQ a -> JQ [(a,Path)]
 collectPaths prog = do
-  arg <- ask
-  return $ runWriterT $ runReaderT prog arg
+  rd <- ask
+  return $ runWriterT $ runReaderT prog rd
 
-insert :: JSValue -> (JSValue, Path) -> JSValue
-insert base (replace, []) = replace
-insert (JSArray values) (replace, ((Left n):rest)) = JSArray values'
-    where
-      (left, (_:right)) = splitAt n values
-      values' = left ++ [replace] ++ right
-insert (JSObject obj) (replace, ((Right k):rest))= JSObject $ toJSObject obj'
-    where
-      withoutK = filter ((/= k) . fst) $ fromJSObject obj
-      obj' = (k, replace):withoutK
+withDefn :: (String, [String], Filter) -> Filter -> Filter
+withDefn (name, formals, body) subexp = 
+  local (\(v,s,d) -> (v,s,M.insert (name,length formals) (Operator func) d)) subexp
+  where
+    func args = local (\(v,s,d) -> (v,M.empty,M.fromList (zip (zip formals (repeat 0)) (map (Operator . const) args)) `M.union` d)) body
+
+subexp :: JQ a -> JQ a
+subexp = censor $ const []
+
+
+yieldPaths :: [(a,Path)] -> JQ a
+yieldPaths ps = ReaderT $ const $ WriterT ps
+
+
+insert :: JQ JSValue -> JSValue -> Path -> JQ JSValue
+insert replace base [] = replace
+insert replace (JSArray values) ((Left n):rest) = do
+  let array = take (max (n+1) (length values)) (values ++ repeat JSNull)
+  replacement <- insert replace (array !! n) rest
+  let (left, (_:right)) = splitAt n array
+  return $ JSArray $ left ++ [replacement] ++ right
+insert replace (JSObject obj) ((Right k):rest) = do
+  let oldval = maybe JSNull id (lookup k $ fromJSObject obj)
+  replacement <- insert replace oldval rest
+  let withoutK = filter ((/= k) . fst) $ fromJSObject obj
+  return $ JSObject $ toJSObject $ (k, replacement):withoutK
+
+insert replace JSNull p@((Right k):rest) = insert replace (JSObject $ toJSObject []) p
+insert replace JSNull p@((Left n):rest)  = insert replace (JSArray []) p
+insert _ base p = error $ "Cannot insert into " ++ intercalate ", " (map (either show show) p) ++ " of " ++ encode base
 
 
 eqj a b = JSBool $ a == b
 
+boolj (JSBool false) = False
+boolj (JSNull) = False
+boolj _ = True
+
+andj a b = JSBool $ boolj a && boolj b
+orj a b = JSBool $ boolj a || boolj b
 
 liftp :: (JSValue -> JSValue) -> JQ JSValue
-liftp f = liftM f ask
+liftp f = liftM f input
 
 idp = undefined
 failp t = []
@@ -68,7 +106,7 @@ anyj values = any isTrue values
 selectp prog = do
   match <- collect prog
   guard $ anyj match
-  ask
+  input
 
 constStr :: String -> JQ JSValue
 constStr = return . JSString . toJSString
@@ -76,9 +114,13 @@ constStr = return . JSString . toJSString
 constInt :: Int -> JQ JSValue
 constInt = return . JSRational False . toRational
 
-updatep p = do
-  t <- ask
-  liftM (foldl insert t) $ collectPaths p
+tr x = trace (show x) x
+
+
+assignp sel replace = do
+  paths <- collectPaths sel
+  t <- input
+  foldM (\base (val,path) -> insert (return val >| replace) base path) t paths
 
 arrayp prog = liftM JSArray $ collect prog
 
@@ -88,15 +130,15 @@ childp' (JSArray values) = msum [tell [Left i] >> return v | (v,i) <- zip values
 childp' (JSObject obj) = msum [tell [Right k] >> return v | (k,v) <- fromJSObject obj]
 childp' _ = mzero
 
-childp = ask >>= childp'
+childp = input >>= childp'
 
 --findp :: Program -> Program
 findp prog = do
   found <- collect prog
-  if anyj found then ask else childp >| findp prog
+  if anyj found then input else childp >| findp prog
 
 groupp prog = do
-  list <- ask
+  list <- input
   case list of
     JSArray values -> do
       marked <- forM values $ \v -> do 
@@ -109,17 +151,72 @@ groupp prog = do
         marked
     _ -> return JSNull
 
+recp prog = do
+  found <- collectPaths prog
+  let (roots,subs) = partition (null . snd) found
+  msum $ 
+    [tell p >> return x | (x,p) <- roots] ++
+    [tell p >> (return x >| recp prog) | (x,p) <- subs]
+
+elsep p1 p2 = do
+  p1' <- collectPaths p1
+  if null p1' then p2 else yieldPaths p1'
+
+fullresultp prog = do
+  res <- collectPaths prog
+  msum [return $ JSObject $ toJSObject $ [("val",a),("path",JSArray $ map fromPath p)] | (a,p) <- res]
+    where
+      fromPath (Left n) = js n
+      fromPath (Right s) = js s
 
-  
 
 withArray f (JSArray values) = JSArray $ f values
 withArray f x = x
 
+withString f (JSString str) = JSString $ toJSString $ f $ fromJSString str
+withString f x = x
+
+{-
 callp "select" [p] = selectp p
 callp "find" [p] = findp p
-callp "set" [p] = updatep p
 callp "sort" [] = liftp (withArray sort)
 callp "group" [p] = groupp p
+callp "rec" [p] = recp p
+callp "empty" [] = mzero
+callp "true" [] = return $ JSBool True
+callp "false" [] = return $ JSBool False
+callp "null" [] = return $ JSNull
+callp "count" [] = liftp countj
+callp "fullresult" [p] = fullresultp p
+callp "uppercase" [] = liftp $ withString $ map toUpper
+callp "lowercase" [] = liftp $ withString $ map toLower
+-}
+
+primitives = M.fromList [((name,arglen),Operator func) | 
+                         (name,arglen,func) <- prim]
+  where
+    prim = [("if",1,\[p] -> selectp p),
+            ("find", 1, \[p] -> findp p),
+            ("group", 1, \[p] -> groupp p),
+            ("rec", 1, \[p] -> recp p), 
+            ("true", 0, const $ return $ JSBool True),
+            ("false", 0, const $ return $ JSBool False),
+            ("null", 0, const $ return $ JSNull),
+            ("count", 0, const $ liftp countj),
+            ("fullresult", 1, \[p] -> fullresultp p),
+            ("zip", 0, const $ liftp zipj),
+            ("keys", 0, const $ liftp keysj)
+            ]
+
+callp :: String -> [Filter] -> Filter
+
+callp name args = do
+  (v,s,d) <- ask
+  runOperator (d ! (name, length args)) args
+
+countj (JSArray v) = js$ length v
+countj (JSObject o) = js$ length $ fromJSObject o
+countj _ = js$ (1::Int)
 
 lookupj :: JSValue -> JSValue -> JQ JSValue
 lookupj (JSArray values) (JSRational _ n) = do
@@ -133,25 +230,75 @@ lookupj (JSObject obj) (JSString s) = do
   case (lookup (fromJSString s) (fromJSObject obj)) of
     Just x -> return x
     Nothing -> return JSNull
+lookupj JSNull (JSRational _ n) = do
+  tell [Left $ round n]
+  return JSNull
+lookupj JSNull (JSString s) = do
+  tell [Right (fromJSString s)]
+  return JSNull
+--lookupj v i = error $ "Cannot get element " ++ encode i ++ " of " ++ encode v
 lookupj _ _ = mzero
 
 
-plusj (JSRational _ n1) (JSRational _ n2) = JSRational True (n1 + n2)
+plusj (JSRational f1 n1) (JSRational f2 n2) = JSRational (f1 || f2) (n1 + n2)
 plusj (JSString s1) (JSString s2) = JSString $ toJSString (fromJSString s1 ++ fromJSString s2)
 plusj (JSArray a1) (JSArray a2) = JSArray $ a1 ++ a2
+plusj (JSObject o1) (JSObject o2) = JSObject $ toJSObject $ o1' ++ fromJSObject o2
+  where
+    newkeys = map fst $ fromJSObject o2
+    o1' = filter (not . (`elem` newkeys) . fst) $ fromJSObject o1
+
+keysj (JSArray v) = js [0..length v - 1]
+keysj (JSObject obj) = js (map fst $ fromJSObject obj)
+keysj _ = JSArray []
 
+zipj jsonValue = result $ tx values ctor
+  where
+    (result, values) = extract packed
+    (packed, ctor) = case jsonValue of
+      JSArray values -> (values, (\vs' -> JSArray $ [v' | Just v' <- vs']))
+      JSObject jsObject -> 
+        let object = fromJSObject jsObject
+            keys = map fst object
+            values = map snd object
+            build vs' = [(k,v') | (k,Just v') <- zip keys vs']
+        in (values, JSObject . toJSObject . build)
+      
+      _ -> error "only arrays and objects may be zipped"
+
+    extract values | all isArray values = (JSArray, [map Just arr | JSArray arr <- values])
+                   | all isObject values = 
+      let objects = [fromJSObject o | JSObject o <- values]
+          keys = S.toList $ S.fromList [k | obj <- objects, (k,_) <- obj]
+          values' :: [[Maybe JSValue]]
+          values' = [[lookup k object | k <- keys] | object <- objects]
+          result r = JSObject $ toJSObject $ zip keys r
+      in (result, values')
+                   | otherwise = error "elements of zipped value must be all objects or all arrays"
+      where
+        isArray (JSArray a) = True
+        isArray _ = False
+        isObject (JSObject o) = True
+        isObject _ = False
+    
+    head' [] = Nothing
+    head' (x:xs) = x
+    tail' [] = []
+    tail' (x:xs) = xs
+    tx values ctor | all null values = []
+                   | otherwise = ctor (map head' values):tx (map tail' values) ctor
 
 js :: JSON a => a -> JSValue
 js = showJSON
 
 index s = do
-  v <- ask
+  v <- input
   lookupj v (js s)
 
 
 dictp progs = do
   liftM (JSObject . toJSObject) $ forM progs $ \(k,v) -> do
-    JSString k' <- k
-    v' <- v
+    JSString k' <- subexp k
+    v' <- subexp v
     return (fromJSString k', v')
 
diff --git a/Lexer.x b/Lexer.x
index 700c69e6..cd3b18c3 100644
--- a/Lexer.x
+++ b/Lexer.x
@@ -7,13 +7,14 @@ import Control.Monad.Error
 
 $digit = 0-9
 $alpha = [a-zA-Z_]
-@reserved = "."|"["|"]"|","|":"|"("|")"|"{"|"}"|"|"|"=="|"+"
+@reserved = "."|"["|"]"|","|":"|"("|")"|"{"|"}"|"|"|"=="|"+"|"="|"$"|"def"|";"|"else"|"and"|"or"|"as"
 @ident = $alpha [$alpha $digit]*
 @string = \" ($printable)* \"
 
 
 tokens :- 
 
+<0> "#" ($printable # [\n\r])* ;
 <0> $white+             ;
 <0> @reserved           { tok TRes }
 <0> @ident              { tok TIdent }
diff --git a/Main.hs b/Main.hs
index 695520cb..9e399d4c 100644
--- a/Main.hs
+++ b/Main.hs
@@ -3,11 +3,11 @@ import Lexer
 import JQ
 import Text.JSON
 import Text.JSON.String
+import PrettyJSON
 import System.Environment
 import Control.Monad
 import System.IO
 
-
 parseJS :: String -> JSValue
 parseJS s = case runGetJSON readJSValue s of
   Left err -> error err
@@ -16,7 +16,8 @@ parseJS s = case runGetJSON readJSValue s of
   
 main = do
   [program] <- getArgs
+  stdlib <- openFile "stdlib.jq" ReadMode >>= hGetContents
   json <- liftM parseJS $ hGetContents stdin
-  case runLexer program >>= runParser of
+  case runLexer (stdlib ++ program) >>= runParser of
     Left err -> putStrLn err
-    Right program -> mapM_ (putStrLn . encode) (runJQ program json)
\ No newline at end of file
+    Right program -> mapM_ (putStrLn . show . renderJSON) (runJQ program json)
\ No newline at end of file
diff --git a/Makefile b/Makefile
new file mode 100644
index 00000000..ba89e9cc
--- /dev/null
+++ b/Makefile
@@ -0,0 +1,8 @@
+jq: *.hs Parser.hs Lexer.hs
+	ghc *.hs -o jq
+
+Parser.hs: Parser.y
+	happy -i Parser.y
+
+Lexer.hs: Lexer.x
+	alex Lexer.x
\ No newline at end of file
diff --git a/Parser.y b/Parser.y
index 544fe5b4..3ea37b05 100644
--- a/Parser.y
+++ b/Parser.y
@@ -7,13 +7,21 @@ import Debug.Trace
 import Data.List
 import Control.Monad.Error
 import Control.Monad.Reader
+
+instance Error (Maybe a) where
+  noMsg = Nothing
+  strMsg = const Nothing
+
+instance (Error a, Error b) => Error (a, b) where
+  noMsg = (noMsg, noMsg)
+  strMsg s = (strMsg s, strMsg s)
 }
 
-%name runParser Exp
+%name runParser TopLevel
 %tokentype { Token }
 
 %monad { Either String }
-%error { \t -> fail $ "parse error: unexpected " ++ if null t then [] else (show $ head t) }
+%error { \t -> throwError $ "parse error: unexpected " ++ if null t then [] else (show $ t) }
 
 %token
   '|' { TRes "|" }
@@ -24,28 +32,61 @@ import Control.Monad.Reader
   '}' { TRes "}" }
   '(' { TRes "(" }
   ')' { TRes ")" }
+  '$' { TRes "$" }
+  'as' { TRes "as" }
   ',' { TRes "," }
   ':' { TRes ":" }
   '==' { TRes "==" }
+  '=' { TRes "=" }
   '+' { TRes "+" }
+  'def' { TRes "def" }
+  ';' { TRes ";" }
+  'else' { TRes "else" }
+  'and' { TRes "and" }
+  'or' { TRes "or" }
   Ident { TIdent $$ }
   String { TString $$ }
   Int { TInt $$ }
 
-%left '|'
+
+%left 'else'
+%right '|'
+%left '='
 %left ','
+%left 'and' 'or'
 %nonassoc '=='
 %left '+'
 
 %%
 
+TopLevel
+  : Defn TopLevel	{ withDefn $1 $2 }
+  | Exp			{ $1 }
+
 Exp 
-  : Exp '|' Exp         { $1 >| $3 }
+  : Exp 'else' Exp      { $1 `elsep` $3 }
+  | Assign '|' Exp	{ do { v <- snd $1; setvar (fst $1) v $3 } }
+  | Exp '|' Exp         { $1 >| $3 }
   | Exp ',' Exp         { $1 `mplus` $3 }
-  | Exp '==' Exp        { liftM2 eqj $1 $3 }
-  | Exp '+' Exp         { liftM2 plusj $1 $3 }
+  | Exp 'and' Exp	{ liftM2 andj $1 $3 }
+  | Exp 'or' Exp	{ liftM2 orj $1 $3 }
+  | Exp '=' Exp         { assignp $1 $3 }
+  | Exp '==' Exp        { liftM2 eqj (subexp $1) (subexp $3) }
+  | Exp '+' Exp         { liftM2 plusj (subexp $1) (subexp $3) }
   | Term                { $1 }
 
+Assign
+  : Term 'as' '$' Ident { ($4, $1) }
+
+Defn
+  : 'def' Ident '=' Exp ';' { ($2, [], $4) }
+  | 'def' Ident '(' ParamList ')' '=' Exp ';' { ($2, $4, $7) }
+
+ParamList
+  :			{ [] }
+  | Ident		{ [$1] }
+  | Ident ';' ParamList { $1:$3 }
+  
 ExpD
   : ExpD '|' ExpD       { $1 >| $3 }
   | ExpD '==' ExpD      { liftM2 eqj $1 $3 }
@@ -53,24 +94,26 @@ ExpD
 
 
 Term 
-  : '.'                 { ask }
+  : '.'                 { input }
   | Term '.' Ident      { $1 >| index $3 }
   | '.' Ident           { index $2 }
   | String              { constStr $1 }
-  | Term '[' Exp ']'    { do {t <- $1; i <- $3; lookupj t i} }
+  | Term '[' Exp ']'    { do {t <- $1; i <- subexp $3; lookupj t i} }
   | Term '[' ']'        { $1 >| childp }
   | '(' Exp ')'         { $2 }
   | '[' Exp ']'         { arrayp $2 }
+  | '[' ']' 		{ arrayp (callp "empty" []) }
   | Int                 { constInt $1 }
   | '{' MkDict '}'      { dictp $2 }
   | Ident '(' Exp ')'   { callp $1 [$3] }
   | Ident               { callp $1 [] }
+  | '$' Ident		{ getvar $2 }
 
 MkDict
   :                     { [] }
   | MkDictPair          { [$1] }
   | MkDictPair ',' MkDict { $1:$3 }
- 
+
 MkDictPair
   : Ident ':' ExpD      { (constStr $1, $3) }
   | Ident               { (constStr $1, index $1) }
diff --git a/PrettyJSON.hs b/PrettyJSON.hs
new file mode 100644
index 00000000..8fe243b3
--- /dev/null
+++ b/PrettyJSON.hs
@@ -0,0 +1,12 @@
+module PrettyJSON where
+import Text.JSON
+import Text.PrettyPrint
+
+renderJSON (JSArray vals) = brackets $ fsep $ punctuate comma $ map renderJSON vals
+
+renderJSON (JSObject jsObject) = 
+  let object = fromJSObject jsObject
+  in braces $ fsep $ punctuate comma $ 
+     [hang (renderJSON (JSString $ toJSString $ k) <> colon) 2 (renderJSON v)
+     | (k,v) <- object]
+renderJSON x = text $ encode x
\ No newline at end of file
diff --git a/stdlib.jq b/stdlib.jq
new file mode 100644
index 00000000..0da7030f
--- /dev/null
+++ b/stdlib.jq
@@ -0,0 +1,11 @@
+def map(f) = [.[] | f];
+
+def first = .[0];
+# def last = .[count-1];
+def next = .[count];
+
+# ([])[] would be a decent definition of "empty"
+# except ([]) is defined as syntactic sugar for empty
+def empty = {}[];
+
+def sort = [group(.) | .[]];
\ No newline at end of file