Add functions

bench/Main.hs

@@ -1,44 +1,45 @@
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+
 import           TinyLang.Field.Generator  ()
 import           TinyLang.Field.Typed.Core
 
+import           Control.Lens
 import           Control.Monad
 import           Data.List
+import           Data.Semigroup
 import           Test.QuickCheck
 
--- A couple of functions for checking the output of generators
-progNodes :: Program f -> Int
-progNodes = stmtsNodes . _programStatements
+progNodes :: Program f -> Sum Int
+progNodes prog =
+    foldMapOf progSubExt (const (Sum 1)) prog
+    <> foldMapOf progSubStatements stmtsNodes prog
 
-stmtsNodes :: Statements f -> Int
-stmtsNodes = sum . map stmtNodes . unStatements
+stmtsNodes :: Statements f -> Sum Int
+stmtsNodes stmts = foldMapOf stmtsSubStatement stmtNodes stmts
 
-stmtNodes :: Statement f -> Int
-stmtNodes (ELet _ e)         = 1 + exprNodes e
-stmtNodes (EAssert e)        = 1 + exprNodes e
+stmtNodes :: Statement f -> Sum Int
+stmtNodes stmt = Sum 1 <> foldMapOf stmtSubExpr exprNodes stmt
 
-exprNodes :: Expr f a -> Int
-exprNodes (EConst _)          = 1
-exprNodes (EVar   _)          = 1
-exprNodes (EAppUnOp _ e)      = 1 + exprNodes e
-exprNodes (EAppBinOp _ e1 e2) = 1 + exprNodes e1 + exprNodes e2
-exprNodes (EIf e e1 e2)       = 1 + exprNodes e + exprNodes e1 + exprNodes e2
+exprNodes :: SomeUniExpr f -> Sum Int
+exprNodes e = Sum 1 <> foldMapOf exprSubExpr exprNodes e
 
-progDepth :: Program f -> Int
-progDepth = stmtsDepth . _programStatements
+-- NOTE:  We need Max 0 for the empty case
+progDepth :: Program f -> Max Int
+progDepth prog =
+    Max 0
+    <> foldMapOf progSubExt (const (Max 1)) prog
+    <> foldMapOf progSubStatements stmtsDepth prog
 
-stmtsDepth :: Statements f -> Int
-stmtsDepth = maximum . (0:) . map stmtDepth . unStatements
+stmtsDepth :: Statements f -> Max Int
+stmtsDepth stmts = Max 0 <> foldMapOf stmtsSubStatement stmtDepth stmts
 
-stmtDepth :: Statement f -> Int
-stmtDepth (ELet _ e)         = 1 + exprDepth e
-stmtDepth (EAssert e)        = 1 + exprDepth e
+stmtDepth :: Statement f -> Max Int
+stmtDepth stmt = (+1) <$> Max 0 <> foldMapOf stmtSubExpr exprDepth stmt
 
-exprDepth :: Expr f a -> Int
-exprDepth (EConst _)          = 1
-exprDepth (EVar _)            = 1
-exprDepth (EAppUnOp _ e)      = 1 + exprDepth e
-exprDepth (EAppBinOp _ e1 e2) = 1 + max (exprDepth e1) (exprDepth e2)
-exprDepth (EIf e e1 e2)       = 1 + max (exprDepth e)  (max (exprDepth e1) (exprDepth e2))
+exprDepth :: SomeUniExpr f -> Max Int
+exprDepth expr = (+1) <$> Max 0 <>  foldMapOf exprSubExpr exprDepth expr
 
 data TestResult = TestResult { nodes :: Int
                              , depth :: Int
@@ -48,7 +49,7 @@ data TestResult = TestResult { nodes :: Int
 runGen :: Int -> IO TestResult
 runGen size = do
     prog <- generate (resize size arbitrary) :: IO (Program (AField Rational))
-    pure $ TestResult (progNodes prog) (progDepth prog)
+    pure $ TestResult (getSum (progNodes prog)) (getMax (progDepth prog))
 
 average :: (Real a, Fractional b) => [a] -> b
 average xs = realToFrac (sum xs) / genericLength xs

field/TinyLang/Field/Core.hs

@@ -4,9 +4,14 @@
 module TinyLang.Field.Core
     ( Program (..)
     , Statements (..)
+    , progSubExt
+    , progSubStatements
+    , progSubStatement
+    , stmtsSubStatement
     ) where
 
 import           Data.Bifunctor
+import           Control.Lens
 import           GHC.Generics
 import           Quiet
 
@@ -16,8 +21,8 @@ newtype Statements stmt = Statements { unStatements :: [stmt] }
     deriving (Show) via (Quiet (Statements stmt))
 
 -- | Basic wrapper of program
-data Program var stmt = Program
-    { _programExts :: [var]
+data Program ext stmt = Program
+    { _programExts :: [ext]
     , _programStatements :: Statements stmt
     }
     deriving (Eq, Foldable, Traversable, Functor)
@@ -26,5 +31,21 @@ instance Bifunctor Program where
     bimap f g (Program exts stmts) = Program (fmap f exts) (fmap g stmts)
 
 -- NOTE:  Adding explicit Show instance to avoid record syntax
-instance (Show var, Show stmt) => Show (Program var stmt) where
+instance (Show ext, Show stmt) => Show (Program ext stmt) where
     show (Program exts stmts) = "Program " ++ show exts ++ " " ++ show stmts
+
+-- Some Traversals
+progSubExt :: Traversal' (Program ext stmts) ext
+progSubExt f = \case
+    Program exts stmts -> Program <$> traverse f exts <*> pure stmts
+
+progSubStatements :: Traversal' (Program ext stmt) (Statements stmt)
+progSubStatements f = \case
+    Program exts stmts -> Program exts <$> f stmts
+
+progSubStatement :: Traversal' (Program ext stmt) stmt
+progSubStatement f = \case
+    Program exts stmts -> Program exts <$> traverse f stmts
+
+stmtsSubStatement :: Traversal' (Statements stmt) stmt
+stmtsSubStatement = traverse

field/TinyLang/Field/Raw/Core.hs

@@ -8,20 +8,32 @@ module TinyLang.Field.Raw.Core
     , UnOp(..)
     , Statement(..)
     , Program
-    , pattern C.Program
+    , pattern Program
     , C._programStatements
+    , C._programExts
     , Statements
-    , pattern C.Statements
+    , pattern Statements
     , C.unStatements
     , RawProgram
     , RawStatements
     , RawStatement
     , RawExpr
+    , progSubExt
+    , progSubStatements
+    , progSubStatement
+    , stmtsSubStatement
+    , stmtSubStatement
+    , progSubExpr
+    , stmtsSubExpr
+    , stmtSubExpr
+    , exprSubExpr
     ) where
 
-import           TinyLang.Field.Uni  hiding (Uni)
+import qualified TinyLang.Field.Uni      as U
+import qualified TinyLang.Field.Type     as T
 import qualified TinyLang.Field.Core     as C
 
+import           Control.Lens
 import           GHC.Generics
 import           Quiet
 
@@ -50,22 +62,34 @@ statement level; the operations acting on statement level are not necessarily
 mappable over a list of statements.
 -}
 
-type Program    v f  = C.Program    (v, SomeUni f) (Statement v f)
-type Statements v f  = C.Statements                (Statement v f)
+-- NOTE: Exts are annotated with SomeUni, as we only support constant externals.
+type Program v f = C.Program (v, U.SomeUni f) (Statement v f)
+
+{-# COMPLETE Program #-}
+pattern Program :: [(v, U.SomeUni f)] -> Statements v f -> Program v f
+pattern Program exts stmts = C.Program exts stmts
+
+
+type Statements v f = C.Statements (Statement v f)
+
+{-# COMPLETE Statements #-}
+pattern Statements :: [Statement v f] -> Statements v f
+pattern Statements stmts = C.Statements stmts
+
 
 data Statement v f
-    = ELet    (v, SomeUni f) (Expr v f)
+    = ELet    (v, T.Type f) (Expr v f)
     | EAssert (Expr v f)
     | EFor    v              Integer    Integer (Statements v f)
     deriving (Show)
 
 data Expr v f
-    = EConst     (SomeUniConst f)
+    = EConst     (U.SomeUniConst f)
     | EVar       v
-    | EAppBinOp  BinOp           (Expr v f) (Expr v f)
-    | EAppUnOp   UnOp            (Expr v f)
-    | EIf        (Expr v f)      (Expr v f) (Expr v f)
-    | ETypeAnn   (SomeUni f)     (Expr v f)
+    | EAppBinOp  BinOp              (Expr v f) (Expr v f)
+    | EAppUnOp   UnOp               (Expr v f)
+    | EIf        (Expr v f)         (Expr v f) (Expr v f)
+    | ETypeAnn   (T.Type f)         (Expr v f)
     deriving (Show)
 
 data BinOp
@@ -98,3 +122,42 @@ type RawProgram    f = Program    Var f
 type RawStatements f = Statements Var f
 type RawStatement  f = Statement  Var f
 type RawExpr       f = Expr       Var f
+
+
+-- Traversals
+progSubExt :: Traversal' (Program v f) (v, U.SomeUni f)
+progSubExt = C.progSubExt
+
+progSubStatements :: Traversal' (Program v f) (Statements v f)
+progSubStatements = C.progSubStatements
+
+progSubStatement :: Traversal' (Program v f) (Statement v f)
+progSubStatement = C.progSubStatement
+
+stmtsSubStatement :: Traversal' (Statements v f) (Statement v f)
+stmtsSubStatement = C.stmtsSubStatement
+
+stmtSubStatement :: Traversal' (Statement v f) (Statement v f)
+stmtSubStatement f = \case
+   EFor var i j stmts -> EFor var i j <$> stmtsSubStatement f stmts
+   x -> pure x
+
+progSubExpr :: Traversal' (Program v f) (Expr v f)
+progSubExpr = progSubStatements . stmtsSubExpr
+
+stmtsSubExpr :: Traversal' (Statements v f) (Expr v f)
+stmtsSubExpr = stmtsSubStatement . stmtSubExpr
+
+stmtSubExpr :: Traversal' (Statement v f) (Expr v f)
+stmtSubExpr f = \case
+    ELet var expr -> ELet var <$> f expr
+    EAssert expr -> EAssert <$> f expr
+    EFor var i j stmts -> EFor var i j <$> stmtsSubExpr f stmts
+
+exprSubExpr :: Traversal' (Expr v f) (Expr v f)
+exprSubExpr f = \case
+    EAppUnOp  unOp  expr1       -> EAppUnOp  unOp  <$> f expr1
+    EAppBinOp binOp expr1 expr2 -> EAppBinOp binOp <$> f expr1 <*> f expr2
+    EIf       expr  expr1 expr2 -> EIf <$> f expr  <*> f expr1 <*> f expr2
+    ETypeAnn  typ   expr        -> ETypeAnn typ    <$> f expr
+    x                           -> pure x