<commit_before>{-# LANGUAGE ExistentialQuantification #-}

module Main where

import Circuit
import Circuit.Conversion
import Circuit.Optimizer
import Circuit.Utils
import qualified Circuit.Format.Acirc    as Acirc
import qualified Circuit.Format.Nigel    as Nigel
import qualified Circuit.Format.Netlist  as Netlist
import qualified Circuit.Format.Graphviz as Graphviz

import qualified Examples.AES             as AES
import qualified Examples.ApplebaumRaykov as AR
import qualified Examples.Comparison      as Comparison
import qualified Examples.Garbler         as Garbler
import qualified Examples.Goldreich       as Goldreich
import qualified Examples.GGM             as GGM
import qualified Examples.Point           as Point
import qualified Examples.Tribes          as Tribes

import Control.Monad
import Lens.Micro.Platform
import System.Exit
import System.FilePath.Posix (takeBaseName, takeExtension)
import Text.Printf
import qualified Data.Map as M
import qualified Data.Text.IO as T

import Options.Applicative
import Data.Semigroup ((<>))

data Mode = CompileAcirc String
          | ReadCircuit FilePath
          | Garble FilePath
          deriving (Show)

data Opts = Opts { mode                :: Mode
                 , verbose             :: Bool
                 , show_info           :: Bool
                 , dont_generate_tests :: Bool
                 , target              :: Maybe String
                 , optimization_level  :: Int
                 , coerce              :: Maybe String
                 } deriving (Show)

parseArgs :: IO Opts
parseArgs = execParser $ info (parser <**> helper)
    (fullDesc <> progDesc "cxs is a tool to compile, convert, and optimize circuits")
  where
    parser = Opts
            <$> subparser
                (command "read"
                    (info (readParser <**> helper)
                        (progDesc "Read an existing circuit"))
                <> command "compile"
                    (info (compileParser <**> helper)
                        (progDesc "Compile a DSL arithmetic circuit"))
                <> command "garble"
                    (info (garbleParser <**> helper)
                        (progDesc "Compile a garler for an existing boolean circuit")))
            <*> switch ( short 'v' <> help "Verbose mode")
            <*> switch ( short 'i' <> help "Print circuit info")
            <*> switch ( short 'T' <> help "Skip test generation")
            <*> (optional $ strOption
                ( short 'o'
                <> metavar "FILE"
                <> help "Write output to file FILE in a known circuit format"))
            <*> option auto
                ( short 'O'
                <> showDefault
                <> value 0
                <> metavar "INT"
                <> help "Optimization level")
            <*> (optional $ strOption
                ( short 't'
                <> metavar "TYPE"
                <> help "Coerce circuit to type TYPE [a,a2,b]"))

    readParser    = ReadCircuit <$> strArgument (metavar "CIRCUIT" <> help "The source circuit")
    compileParser = CompileAcirc <$> (switch (short 'c' <> help "Compile a DSL circuit from the examples")
                        *> strArgument (metavar "NAME" <> help "The name of the compilation routine"))
    garbleParser  = Garble <$> strArgument (metavar "CIRCUIT" <> help "The source circuit to garble (must be boolean)")

main :: IO ()
main = chooseMode =<< parseArgs

chooseMode :: Opts -> IO ()
chooseMode opts = do
    when (verbose opts) (print opts)
    case mode opts of
        CompileAcirc name -> do
            let m = M.union
                        (include [Point.export, Garbler.export, GGM.export, Goldreich.export]) $
                        M.fromList
                        [ ("aes"           , compile opts AES.make)
                        , ("aes1r"         , compile opts AES.makeAes1r)
                        , ("aes10r"        , compile opts AES.makeAes10r)
                        , ("applebaum"     , compile opts AR.makeApplebaum)
                        , ("tribes"        , compile opts Tribes.make)
                        , ("comparison"    , compile opts Comparison.make)
                        ]
            case M.lookup name m of
                Just c  -> c
                Nothing -> do
                    printf "[main] unknown circuit generation mode \"%s\"!\nknown modes:\n" name
                    mapM_ (printf "\t%s\n") (M.keys m)
                    exitFailure
        ReadCircuit inp -> do
            let ext = takeExtension inp

            when (ext `notElem` [".acirc", ".nigel", ".netlist"]) $
                error (printf "[main] unknown input extension \"%s\"!" ext)

            case target opts of
                Nothing -> case ext of
                    ".acirc" -> case coerce opts of
                        Just "a"  -> run opts (Acirc.readWithTests inp)
                        Just "a2" -> run opts (over _1 toAcirc2 <$> Acirc.readWithTests inp)
                        Just _    -> error "[main] supported types for acirc: [a,a2]"
                        Nothing   -> run opts (Acirc.readWithTests inp)
                    ".nigel" -> case coerce opts of
                        Just "a"   -> run opts (Nigel.readNigel inp :: IO (Acirc, [TestCase]))
                        Just "a2"  -> run opts (Nigel.readNigel inp :: IO (Acirc2, [TestCase]))
                        Just "b"   -> run opts (Nigel.readNigel inp :: IO (Acirc2, [TestCase]))
                        Nothing    -> run opts (Nigel.readNigel inp :: IO (Circ, [TestCase]))
                        Just other -> error "[main] supported coersion types for nigel: [a, a2, b]"
                    ".netlist" -> case coerce opts of
                        Just "a"   -> run opts (Netlist.readNetlist inp :: IO (Acirc, [TestCase]))
                        Just "a2"  -> run opts (Netlist.readNetlist inp :: IO (Acirc2, [TestCase]))
                        Just "b"   -> run opts (Netlist.readNetlist inp :: IO (Circ, [TestCase]))
                        Nothing    -> run opts (Netlist.readNetlist inp :: IO (Circ, [TestCase]))
                        Just other -> error "[main] supported coersion types for netlist: [a, a2, b]"

                Just outputFile -> case takeExtension outputFile of
                    ".acirc" -> case ext of
                        ".acirc"   -> run opts (Acirc.readWithTests inp)
                        ".nigel"   -> run opts (Nigel.readNigel inp :: IO (Acirc, [TestCase]))
                        ".netlist" -> run opts (Netlist.readNetlist inp :: IO (Acirc, [TestCase]))
                    ".nigel" -> case ext of
                        ".nigel"   -> run opts (Nigel.readNigel inp :: IO (Circ, [TestCase]))
                        ".netlist" -> run opts (Netlist.readNetlist inp :: IO (Circ, [TestCase]))
                        other      -> error "[main] supported input formats for nigel output: [nigel, netlist]"
                    ".netlist" -> case ext of
                        ".nigel"   -> run opts (Nigel.readNigel inp :: IO (Circ, [TestCase]))
                        ".netlist" -> run opts (Netlist.readNetlist inp :: IO (Circ, [TestCase]))
                        other      -> error "[main] supported input formats for netlist output: [nigel, netlist]"

        Garble inp -> do
            let ext = takeExtension inp
            when (ext `notElem` [".nigel", ".netlist"]) $
                error (printf "[main] \"%s\" is not a known boolean circuit type!" ext)
            c <- case ext of
                ".nigel"   -> Nigel.read inp
                ".netlist" -> Netlist.read inp
            g <- Garbler.garbler c
            circuitMain opts (Just (printf "garbled_%s.acirc" (takeBaseName inp))) g []

  where
    include = M.unions . map (fmap (compile opts) . M.fromList)
    compile :: Opts -> [(String, IO Acirc)] -> IO ()
    compile opts tups = forM_ tups $ \(fname, comp) -> do
        c <- comp
        circuitMain opts (Just fname) c []
    run opts comp = uncurry (circuitMain opts (target opts)) =<< comp

circuitMain :: (Graphviz.Graphviz g, Optimize g, Gate g, ToAcirc g, ToCirc g, ToAcirc2 g)
            => Opts -> Maybe FilePath -> Circuit g -> [TestCase] -> IO ()
circuitMain opts outputName c ts = do
    let old_symlen = _circ_symlen c

    c' <- case optimization_level opts of
        0 -> return c
        1 -> optimizeO1 c
        2 -> optimizeO2 c
        3 -> optimizeO3 c
        x -> do
            printf "[error] unknown optimization level %d\n" x
            exitFailure

    let c = c' & circ_symlen .~ old_symlen

    when (show_info opts) $ do
        printCircInfo c

    ts <- if dont_generate_tests opts
             then return ts
             else replicateM 10 (genTest c)

    case outputName of
        Nothing -> return ()
        Just fn -> do
            let t = case takeExtension fn of
                    ".acirc" -> Acirc.showWithTests (toAcirc c) ts
                    ".nigel" -> Nigel.showCirc (toCirc c)
                    ".dot"   -> Graphviz.showGraphviz c
                    other    -> error (printf "[main] unknown output format \"%s\"" other)
            printf "writing %s\n" fn
            T.writeFile fn t


evalTests :: Gate g => Opts -> Circuit g -> [TestCase] -> IO ()<commit_msg>Option for number of tests<commit_after>{-# LANGUAGE ExistentialQuantification #-}

module Main where

import Circuit
import Circuit.Conversion
import Circuit.Optimizer
import Circuit.Utils
import qualified Circuit.Format.Acirc    as Acirc
import qualified Circuit.Format.Nigel    as Nigel
import qualified Circuit.Format.Netlist  as Netlist
import qualified Circuit.Format.Graphviz as Graphviz

import qualified Examples.AES             as AES
import qualified Examples.ApplebaumRaykov as AR
import qualified Examples.Comparison      as Comparison
import qualified Examples.Garbler         as Garbler
import qualified Examples.Goldreich       as Goldreich
import qualified Examples.GGM             as GGM
import qualified Examples.Point           as Point
import qualified Examples.Tribes          as Tribes

import Control.Monad
import Lens.Micro.Platform
import System.Exit
import System.FilePath.Posix (takeBaseName, takeExtension)
import Text.Printf
import qualified Data.Map as M
import qualified Data.Text.IO as T

import Options.Applicative
import Data.Semigroup ((<>))

data Mode = CompileAcirc String
          | ReadCircuit FilePath
          | Garble FilePath
          deriving (Show)

data Opts = Opts { mode                :: Mode
                 , verbose             :: Bool
                 , show_info           :: Bool
                 , ntests              :: Maybe Int
                 , target              :: Maybe String
                 , optimization_level  :: Int
                 , coerce              :: Maybe String
                 } deriving (Show)

parseArgs :: IO Opts
parseArgs = execParser $ info (parser <**> helper)
    (fullDesc <> progDesc "cxs is a tool to compile, convert, and optimize circuits")
  where
    parser = Opts
            <$> subparser
                (command "read"
                    (info (readParser <**> helper)
                        (progDesc "Read an existing circuit"))
                <> command "compile"
                    (info (compileParser <**> helper)
                        (progDesc "Compile a DSL arithmetic circuit"))
                <> command "garble"
                    (info (garbleParser <**> helper)
                        (progDesc "Compile a garler for an existing boolean circuit")))
            <*> switch (short 'v' <> help "Verbose mode")
            <*> switch (short 'i' <> help "Print circuit info")
            <*> (optional $ option auto (short 'T' <> metavar "N" <> help "Generate N tests"))
            <*> (optional $ strOption
                ( short 'o'
                <> metavar "FILE"
                <> help "Write output to file FILE in a known circuit format"))
            <*> option auto
                ( short 'O'
                <> showDefault
                <> value 0
                <> metavar "N"
                <> help "Optimization level")
            <*> (optional $ strOption
                ( short 't'
                <> metavar "TYPE"
                <> help "Coerce circuit to type TYPE [a,a2,b]"))

    readParser    = ReadCircuit <$> strArgument (metavar "CIRCUIT" <> help "The source circuit")
    compileParser = CompileAcirc <$> (switch (short 'c' <> help "Compile a DSL circuit from the examples")
                        *> strArgument (metavar "NAME" <> help "The name of the compilation routine"))
    garbleParser  = Garble <$> strArgument (metavar "CIRCUIT" <> help "The source circuit to garble (must be boolean)")

main :: IO ()
main = chooseMode =<< parseArgs

chooseMode :: Opts -> IO ()
chooseMode opts = do
    when (verbose opts) (print opts)
    case mode opts of
        CompileAcirc name -> do
            let m = M.union
                        (include [Point.export, Garbler.export, GGM.export, Goldreich.export]) $
                        M.fromList
                        [ ("aes"           , compile opts AES.make)
                        , ("aes1r"         , compile opts AES.makeAes1r)
                        , ("aes10r"        , compile opts AES.makeAes10r)
                        , ("applebaum"     , compile opts AR.makeApplebaum)
                        , ("tribes"        , compile opts Tribes.make)
                        , ("comparison"    , compile opts Comparison.make)
                        ]
            case M.lookup name m of
                Just c  -> c
                Nothing -> do
                    printf "[main] unknown circuit generation mode \"%s\"!\nknown modes:\n" name
                    mapM_ (printf "\t%s\n") (M.keys m)
                    exitFailure
        ReadCircuit inp -> do
            let ext = takeExtension inp

            when (ext `notElem` [".acirc", ".nigel", ".netlist"]) $
                error (printf "[main] unknown input extension \"%s\"!" ext)

            case target opts of
                Nothing -> case ext of
                    ".acirc" -> case coerce opts of
                        Just "a"  -> run opts (Acirc.readWithTests inp)
                        Just "a2" -> run opts (over _1 toAcirc2 <$> Acirc.readWithTests inp)
                        Just _    -> error "[main] supported types for acirc: [a,a2]"
                        Nothing   -> run opts (Acirc.readWithTests inp)
                    ".nigel" -> case coerce opts of
                        Just "a"   -> run opts (Nigel.readNigel inp :: IO (Acirc, [TestCase]))
                        Just "a2"  -> run opts (Nigel.readNigel inp :: IO (Acirc2, [TestCase]))
                        Just "b"   -> run opts (Nigel.readNigel inp :: IO (Acirc2, [TestCase]))
                        Nothing    -> run opts (Nigel.readNigel inp :: IO (Circ, [TestCase]))
                        Just other -> error "[main] supported coersion types for nigel: [a, a2, b]"
                    ".netlist" -> case coerce opts of
                        Just "a"   -> run opts (Netlist.readNetlist inp :: IO (Acirc, [TestCase]))
                        Just "a2"  -> run opts (Netlist.readNetlist inp :: IO (Acirc2, [TestCase]))
                        Just "b"   -> run opts (Netlist.readNetlist inp :: IO (Circ, [TestCase]))
                        Nothing    -> run opts (Netlist.readNetlist inp :: IO (Circ, [TestCase]))
                        Just other -> error "[main] supported coersion types for netlist: [a, a2, b]"

                Just outputFile -> case takeExtension outputFile of
                    ".acirc" -> case ext of
                        ".acirc"   -> run opts (Acirc.readWithTests inp)
                        ".nigel"   -> run opts (Nigel.readNigel inp :: IO (Acirc, [TestCase]))
                        ".netlist" -> run opts (Netlist.readNetlist inp :: IO (Acirc, [TestCase]))
                    ".nigel" -> case ext of
                        ".nigel"   -> run opts (Nigel.readNigel inp :: IO (Circ, [TestCase]))
                        ".netlist" -> run opts (Netlist.readNetlist inp :: IO (Circ, [TestCase]))
                        other      -> error "[main] supported input formats for nigel output: [nigel, netlist]"
                    ".netlist" -> case ext of
                        ".nigel"   -> run opts (Nigel.readNigel inp :: IO (Circ, [TestCase]))
                        ".netlist" -> run opts (Netlist.readNetlist inp :: IO (Circ, [TestCase]))
                        other      -> error "[main] supported input formats for netlist output: [nigel, netlist]"

        Garble inp -> do
            let ext = takeExtension inp
            when (ext `notElem` [".nigel", ".netlist"]) $
                error (printf "[main] \"%s\" is not a known boolean circuit type!" ext)
            c <- case ext of
                ".nigel"   -> Nigel.read inp
                ".netlist" -> Netlist.read inp
            g <- Garbler.garbler c
            circuitMain opts (Just (printf "garbled_%s.acirc" (takeBaseName inp))) g []

  where
    include = M.unions . map (fmap (compile opts) . M.fromList)
    compile :: Opts -> [(String, IO Acirc)] -> IO ()
    compile opts tups = forM_ tups $ \(fname, comp) -> do
        c <- comp
        circuitMain opts (Just fname) c []
    run opts comp = uncurry (circuitMain opts (target opts)) =<< comp

circuitMain :: (Graphviz.Graphviz g, Optimize g, Gate g, ToAcirc g, ToCirc g, ToAcirc2 g)
            => Opts -> Maybe FilePath -> Circuit g -> [TestCase] -> IO ()
circuitMain opts outputName c ts = do
    let old_symlen = _circ_symlen c

    c' <- case optimization_level opts of
        0 -> return c
        1 -> optimizeO1 c
        2 -> optimizeO2 c
        3 -> optimizeO3 c
        x -> do
            printf "[error] unknown optimization level %d\n" x
            exitFailure

    let c = c' & circ_symlen .~ old_symlen

    when (show_info opts) $ do
        printCircInfo c

    ts <- case ntests opts of
        Just n  -> replicateM n (genTest c)
        Nothing -> return ts

    case outputName of
        Nothing -> return ()
        Just fn -> do
            let t = case takeExtension fn of
                    ".acirc" -> Acirc.showWithTests (toAcirc c) ts
                    ".nigel" -> Nigel.showCirc (toCirc c)
                    ".dot"   -> Graphviz.showGraphviz c
                    other    -> error (printf "[main] unknown output format \"%s\"" other)
            printf "writing %s\n" fn
            T.writeFile fn t


evalTests :: Gate g => Opts -> Circuit g -> [TestCase] -> IO ()
{'fraction_non_alphanumeric': 0.10389123802387197, 'fraction_numerical': 0.0029973772948669914, 'mean_word_length': 1.5570001368550705, 'pattern_counts': {'":': 0, '<': 90, '<?xml version=': 0, '>': 200, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2035600', 'n_tokens_mistral': 5321, 'n_tokens_neox': 4890, 'n_words': 1713}