diff --git a/ChangeLog.md b/ChangeLog.md new file mode 100644 --- /dev/null +++ b/ChangeLog.md @@ -0,0 +1,30 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [v0.9.0] + +### Changed + +* Search for `aurora_prove` and `aurora_verify` in the environment instead. + +### Removed + +* Temporarily removed Docker support. + +## [v0.8.4] + +### Added + +* `Keelung.Syntax.Monad` + * Functions for requesting fresh varibles: `freshVar`, `freshVarField`, `freshVarBool`, and `freshVarUInt`. + +### Removed + +* `Keelung.Syntax` + * Function `uintToBool` + * Function `fieldToBool` + diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2022-2023 BTQ AG + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,17 @@ +# Keelung + +Keelung is a domain-specific language (DSL) embedded in Haskell that is designed for creating zero-knowledge proofs (ZKPs). With Keelung, developers can create high-quality and reliable ZKPs with ease, even without expertise in cryptography. + +keelung5 + +## Documentation + +Check out our [GitBook](https://btq.gitbook.io/keelung/) for more information. + +## Language Reference + +The language reference will be available on Hackage soon. + +## Standard Library + +The standard library is available [here](https://github.com/btq-ag/keelung-stdlib). diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/keelung.cabal b/keelung.cabal new file mode 100644 --- /dev/null +++ b/keelung.cabal @@ -0,0 +1,71 @@ +cabal-version: 1.12 + +-- This file has been generated from package.yaml by hpack version 0.35.0. +-- +-- see: https://github.com/sol/hpack + +name: keelung +version: 0.9.0.0 +synopsis: DSL for creating zero-knowledge proofs +description: Please see the README on GitHub at +category: Cryptography +homepage: https://github.com/btq-ag/keelung#readme +bug-reports: https://github.com/btq-ag/keelung/issues +author: BTQ AG +maintainer: BTQ AG +copyright: 2022-2023 BTQ AG +license: Apache-2.0 +license-file: LICENSE +build-type: Simple +extra-source-files: + README.md + ChangeLog.md + +source-repository head + type: git + location: https://github.com/btq-ag/keelung + +library + exposed-modules: + Keelung + Keelung.Constraint.R1C + Keelung.Constraint.R1CS + Keelung.Data.BinRep + Keelung.Data.Bits + Keelung.Data.N + Keelung.Data.Polynomial + Keelung.Data.Struct + Keelung.Error + Keelung.Field + Keelung.Heap + Keelung.Monad + Keelung.Prof + Keelung.Syntax + Keelung.Syntax.Counters + Keelung.Syntax.Encode + Keelung.Syntax.Encode.Syntax + other-modules: + Paths_keelung + hs-source-dirs: + src + default-extensions: + ImportQualifiedPost + ghc-options: -Wall -Werror=incomplete-patterns -fno-warn-orphans + build-depends: + QuickCheck + , array + , base >=4.7 && <5 + , bytestring + , cereal + , containers + , deepseq + , directory + , galois-field + , groups + , mtl + , process + , random + , semirings + , vector + , wl-pprint-text + default-language: Haskell2010 diff --git a/src/Keelung.hs b/src/Keelung.hs new file mode 100644 --- /dev/null +++ b/src/Keelung.hs @@ -0,0 +1,350 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | Keelung is a DSL for building zero-knowledge proofs +module Keelung + ( module Keelung.Syntax, + module Keelung.Field, + module Keelung.Heap, + module Keelung.Monad, + module Keelung.Data.Bits, + run, + compile, + compileO0, + compileO2, + compileWithOpts, + rtsoptProf, + rtsoptMemory, + generate, + verify, + genCircuit, + genWitness, + genInputs, + interpret_, + interpret, + gf181, + bn128, + b64, + elaborateAndEncode, + Encode, + GaloisField, + keelungVersion, + ) +where + +import Control.Monad.Except +import Data.ByteString.Char8 qualified as BS +import Data.Field.Galois (GaloisField) +import Data.List (intercalate) +import Data.Serialize (Serialize) +import Data.Serialize qualified as Serialize +import Data.String (IsString (fromString)) +import Keelung.Constraint.R1CS (R1CS) +import Keelung.Data.Bits +import Keelung.Data.Struct (Struct (..)) +import Keelung.Error +import Keelung.Field +import Keelung.Heap +import Keelung.Monad +import Keelung.Syntax +import Keelung.Syntax.Encode +import Keelung.Syntax.Encode.Syntax qualified as Encoding +import System.Directory qualified as Path +import System.IO.Error qualified as IO +import System.Info qualified +import System.Process qualified as Process +import Text.Read (readMaybe) + +-------------------------------------------------------------------------------- + +-- | Compile a program to a 'R1CS' constraint system. +compile :: Encode t => FieldType -> Comp t -> IO (Either Error (R1CS Integer)) +compile = compileWithOpts 1 [] [] + +-- | Compile a program to a 'R1CS' constraint system with optimization level 0. +compileO0 :: Encode t => FieldType -> Comp t -> IO (Either Error (R1CS Integer)) +compileO0 = compileWithOpts 0 [] [] + +-- | Compile a program to a 'R1CS' constraint system with optimization level 2. +compileO2 :: Encode t => FieldType -> Comp t -> IO (Either Error (R1CS Integer)) +compileO2 = compileWithOpts 2 [] [] + +-- | Compile a program to a 'R1CS' constraint system with optimization level and RTS options as arguments. +compileWithOpts :: Encode t => Int -> [String] -> [String] -> FieldType -> Comp t -> IO (Either Error (R1CS Integer)) +compileWithOpts level opts rtsopts fieldType prog = runM $ do + elab <- liftEither (elaborateAndEncode prog) + let opts' = "protocol" : optOptimize level : opts <> ["+RTS"] <> rtsopts <> ["-RTS"] + case fieldType of + GF181 -> convertFieldElement (wrapper opts' (fieldType, elab) :: M (R1CS GF181)) + BN128 -> convertFieldElement (wrapper opts' (fieldType, elab) :: M (R1CS BN128)) + B64 -> convertFieldElement (wrapper opts' (fieldType, elab) :: M (R1CS B64)) + where + optOptimize :: Int -> String + optOptimize i = "O" <> show i + +-- | Default RTS options for profiling +rtsoptProf :: [String] +rtsoptProf = ["-p"] + +-- | Helper function for compiling RTS options +-- Memory size in GB for RTS options -M, -H and in MB for -A +-- Try to increase if keelungc produces segmentation fault. +-- https://downloads.haskell.org/ghc/latest/docs/users_guide/runtime_control.html +rtsoptMemory :: Int -> Int -> Int -> [String] +rtsoptMemory m h a = ["-M" <> show m <> "G", "-H" <> show h <> "G", "-A" <> show a <> "M"] + +-------------------------------------------------------------------------------- + +-- | Generate a proof +generate_ :: (Serialize n, Integral n, Encode t) => FieldType -> Comp t -> [n] -> [n] -> IO (Either Error (FilePath, String)) +generate_ fieldType prog publicInput privateInput = runM $ do + exists <- checkCmd "aurora_prove" + _ <- genCircuit fieldType prog + _ <- genWitness_ fieldType prog publicInput privateInput + proofPath <- lift $ Path.makeAbsolute "proof" + genParameters + if exists + then lift $ do + let arguments = + [ "--r1cs_filepath", + "circuit.jsonl", + "--input_filepath", + "witness.jsonl", + "--parameter_filepath", + "parameter.json", + "--output_filepath", + proofPath + ] + msg <- Process.readProcess "aurora_prove" arguments mempty + return (proofPath, msg) + else throwError CannotLocateProver + +-- | Generate a proof +generate :: Encode t => FieldType -> Comp t -> [Integer] -> [Integer] -> IO () +generate fieldType prog publicInput privateInput = do + result <- generate_ fieldType prog publicInput privateInput + case result of + Left err -> print err + Right (_, msg) -> putStr msg + +-- | Generate and verify a proof +verify_ :: IO (Either Error String) +verify_ = runM $ do + exists <- checkCmd "aurora_verify" + genParameters + if exists + then lift $ do + let arguments = + [ "--r1cs_filepath", + "circuit.jsonl", + "--input_filepath", + "inputs.jsonl", + "--parameter_filepath", + "parameter.json", + "--proof_filepath", + "proof" + ] + Process.readProcess "aurora_verify" arguments mempty + else throwError CannotLocateVerifier + +-- | Verify a proof +verify :: IO () +verify = do + result <- verify_ + case result of + Left err -> print err + Right msg -> putStr msg + +-- | Compile a program as R1CS and write it to circuit.jsonl. +genCircuit :: Encode t => FieldType -> Comp t -> M (R1CS Integer) +genCircuit fieldType prog = do + elab <- liftEither (elaborateAndEncode prog) + case fieldType of + GF181 -> convertFieldElement (wrapper ["protocol", "toJSON"] (fieldType, elab) :: M (R1CS GF181)) + BN128 -> convertFieldElement (wrapper ["protocol", "toJSON"] (fieldType, elab) :: M (R1CS BN128)) + B64 -> convertFieldElement (wrapper ["protocol", "toJSON"] (fieldType, elab) :: M (R1CS B64)) + +-- | Generate witnesses for a program with inputs and write them to witness.jsonl. +genWitness_ :: (Serialize n, Integral n, Encode t) => FieldType -> Comp t -> [n] -> [n] -> M [n] +genWitness_ fieldType prog publicInput privateInput = do + elab <- liftEither (elaborateAndEncode prog) + wrapper ["protocol", "genWitness"] (fieldType, elab, map toInteger publicInput, map toInteger privateInput) + +-- | Generate parameters for a program and write them to parameter.json. +genParameters :: M () +genParameters = lift $ BS.writeFile "parameter.json" "{\"security_level\": 128, \"heuristic_ldt_reducer_soundness\": true, \"heuristic_fri_soundness\": true, \"bcs_hash_type\": \"blake2b_type\", \"make_zk\": false, \"parallel\": true, \"field_size\": 181, \"is_multiplicative\": true}" + +-- | For generating witness.jsonl +genWitness :: Encode t => FieldType -> Comp t -> [Integer] -> [Integer] -> IO [Integer] +genWitness fieldType prog publicInput privateInput = runM (genWitness_ fieldType prog publicInput privateInput) >>= printErrorInstead + +-- | For generating inputs.jsonl +genInputs :: (Integral n) => [n] -> M () +genInputs inputs = do + let inputs' = intercalate "," $ map ((\x -> "\"" ++ x ++ "\"") . show . toInteger) inputs + lift $ BS.writeFile "inputs.jsonl" $ fromString $ "{\"inputs\":[" ++ inputs' ++ "]}" + +-------------------------------------------------------------------------------- + +-- | Interpret a program +interpret_ :: (Serialize n, Integral n, Encode t) => FieldType -> Comp t -> [n] -> [n] -> IO (Either Error [n]) +interpret_ fieldType prog publicInput privateInput = runM $ do + elab <- liftEither (elaborateAndEncode prog) + wrapper ["protocol", "interpret"] (fieldType, elab, map toInteger publicInput, map toInteger privateInput) + +printErrorInstead :: Show e => Either e [a] -> IO [a] +printErrorInstead (Left err) = do + print err + return [] +printErrorInstead (Right values) = return values + +-- | Interpret a program with public and private inputs +run :: Encode t => Comp t -> [Integer] -> [Integer] -> IO [Integer] +run prog publicInput privateInput = interpret_ GF181 prog publicInput privateInput >>= printErrorInstead + +-- | Interpret a program with public and private inputs +interpret :: Encode t => FieldType -> Comp t -> [Integer] -> [Integer] -> IO [Integer] +interpret fieldType prog publicInput privateInput = interpret_ fieldType prog publicInput privateInput >>= printErrorInstead + +-- | A specialized version of 'interpret' that outputs numbers as 'N GF181' +gf181 :: Encode t => Comp t -> [GF181] -> [GF181] -> IO [N GF181] +gf181 prog publicInput privateInput = map N <$> (interpret_ GF181 prog publicInput privateInput >>= printErrorInstead) + +-- | A specialized version of 'interpret' that outputs numbers as 'N B64' +b64 :: Encode t => Comp t -> [B64] -> [B64] -> IO [N B64] +b64 prog publicInput privateInput = map N <$> (interpret_ B64 prog publicInput privateInput >>= printErrorInstead) + +-- | A specialized version of 'interpret' that outputs numbers as 'N BN128' +bn128 :: Encode t => Comp t -> [BN128] -> [BN128] -> IO [N BN128] +bn128 prog publicInput privateInput = map N <$> (interpret_ BN128 prog publicInput privateInput >>= printErrorInstead) + +-------------------------------------------------------------------------------- + +-- | Elaborate a program and encode it +elaborateAndEncode :: Encode t => Comp t -> Either Error Encoding.Elaborated +elaborateAndEncode prog = encodeElaborated <$> elaborate prog + where + encodeElaborated :: Encode t => Elaborated t -> Encoding.Elaborated + encodeElaborated (Elaborated expr comp) = runHeapM (compHeap comp) $ do + let Computation counters _addrSize _heap eb assertions divModRelsU = comp + in Encoding.Elaborated + <$> encode expr + <*> ( Encoding.Computation + counters + <$> (Struct <$> mapM encode' (structF eb) <*> mapM encode' (structB eb) <*> pure (structU eb)) + <*> mapM encode assertions + <*> pure divModRelsU + ) + +-------------------------------------------------------------------------------- + +-- | Internal function for handling data serialization +wrapper :: (Serialize a, Serialize b) => [String] -> a -> M b +wrapper args' payload = do + (cmd, args) <- findKeelungc + version <- readKeelungVersion cmd args + checkKeelungVersion version + blob <- lift $ Process.readProcess cmd (args ++ args') (BS.unpack $ Serialize.encode payload) + let result = Serialize.decode (BS.pack blob) + case result of + Left err -> throwError (DecodeError err) + Right (Left err) -> throwError (CompileError err) + Right (Right x) -> return x + +-- | Locate the Keelung compiler +-- 1. see if "keelungc" is in PATH +-- 2. if not, try to run "docker run banacorn/keelung" +-- Returns the command and arguments to run when found +findKeelungc :: M (String, [String]) +findKeelungc = do + keelungcExists <- checkCmd "keelungc" + if keelungcExists + then return ("keelungc", []) + else do + -- dockerExists <- checkCmd "docker" + let dockerExists = False + if dockerExists + then -- insert "--platform=linux/amd64" when we are not on a x86 machine + case System.Info.arch of + "x86_64" -> return ("docker", ["run", "-i", "banacorn/keelung"]) + _ -> return ("docker", ["run", "-i", "--platform=linux/amd64", "banacorn/keelung"]) + else throwError CannotLocateKeelungC + +-- | Check the version of the Keelung compiler +readKeelungVersion :: FilePath -> [String] -> M (Int, Int, Int) +readKeelungVersion cmd args = do + -- trying to read version with `keelungc --version` + rawString <- + catchIOError + CannotReadVersionError + (Process.readProcess cmd (args ++ ["--version"]) mempty) + -- parse see if the version number is well-formed + let parseResult = case splitAt 9 rawString of + ("Keelung v", versionString) -> parseVersion versionString + _ -> Nothing + -- throws CannotReadVersionError if it's not well-formed + case parseResult of + Nothing -> throwError CannotReadVersionError + Just x -> return x + where + parseVersion :: String -> Maybe (Int, Int, Int) + parseVersion versionString = do + (major, minor, patch) <- case span (/= '.') versionString of + (m, '.' : rest) -> case span (/= '.') rest of + (n, '.' : p) -> Just (m, n, p) + _ -> Nothing + _ -> Nothing + (,,) <$> readMaybe major <*> readMaybe minor <*> readMaybe patch + +checkKeelungVersion :: (Int, Int, Int) -> M () +checkKeelungVersion (major, minor, patch) = do + if major == 0 && minor >= 9 && minor < 10 && patch >= 0 + then return () + else throwError (VersionMismatchError major minor patch) + +-------------------------------------------------------------------------------- + +-- | Check if a command exists +checkCmd :: String -> M Bool +checkCmd cmd = + lift $ + IO.catchIOError + (Process.readProcess whichCmd [cmd] mempty >> return True) + (\_ -> return False) + where + -- decide the command for locating executables + whichCmd :: String + whichCmd = case System.Info.os of + "mingw32" -> "where" -- Windows uses "where" + _ -> "which" -- Unix uses "which" + +-------------------------------------------------------------------------------- + +-- | The version of Keelung is a triple of three numbers, we're not going full semver yet +keelungVersion_ :: (Int, Int, Int) +keelungVersion_ = (0, 9, 0) + +-- | String of Keelung version exposed to the user +keelungVersion :: String +keelungVersion = let (major, minor, patch) = keelungVersion_ in show major ++ "." ++ show minor ++ "." ++ show patch + +-------------------------------------------------------------------------------- + +type M = ExceptT Error IO + +runM :: M a -> IO (Either Error a) +runM = runExceptT + +-- liftEitherT :: IO (Either Error a) -> M a +-- liftEitherT f = do +-- result <- lift f +-- case result of +-- Left err -> throwError err +-- Right x -> return x + +-- | Handle 'IO' Exceptions in the 'M' Monad +catchIOError :: Error -> IO a -> M a +catchIOError err f = lift (IO.catchIOError (Right <$> f) (const (return (Left err)))) >>= liftEither + +-- | Prettify and convert all field elements to 'Integer' in a 'R1CS' +convertFieldElement :: (GaloisField a, Integral a) => M (R1CS a) -> M (R1CS Integer) +convertFieldElement = fmap (fmap (toInteger . N)) diff --git a/src/Keelung/Constraint/R1C.hs b/src/Keelung/Constraint/R1C.hs new file mode 100644 --- /dev/null +++ b/src/Keelung/Constraint/R1C.hs @@ -0,0 +1,111 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} + +-- | Rank-1 Constraint +module Keelung.Constraint.R1C (R1C (..), satisfy, freeVars) where + +import Control.DeepSeq (NFData) +import Data.Field.Galois (GaloisField) +import Data.IntMap (IntMap) +import Data.IntMap qualified as IntMap +import Data.IntSet (IntSet) +import Data.Serialize (Serialize) +import GHC.Generics (Generic) +import Keelung.Data.Polynomial (Poly) +import Keelung.Data.Polynomial qualified as Poly + +-------------------------------------------------------------------------------- + +-- | A Rank-1 Constraint is a relation between 3 polynomials +-- Ax * Bx = Cx +data R1C n = R1C (Either n (Poly n)) (Either n (Poly n)) (Either n (Poly n)) + deriving (Generic, NFData) + +instance (Eq n, Num n) => Eq (R1C n) where + R1C a b c == R1C a' b' c' = + -- if the RHS are the same + (c == c' && (a == a' && b == b' || a == b' && b == a')) + -- if the RHS are the negation of each other + || ( negate' c == c' + && ( negate' a == a' && b == b' + || a == a' && negate' b == b' + || negate' a == b' && b == a' + || a == b' && negate' b == a' + ) + ) + where + negate' (Left n) = Left (-n) + negate' (Right p) = Right (Poly.negate p) + +instance Functor R1C where + fmap f (R1C a b c) = R1C (fmapE a) (fmapE b) (fmapE c) + where + fmapE (Left x) = Left (f x) + fmapE (Right xs) = Right (fmap f xs) + +instance Serialize n => Serialize (R1C n) + +instance (Num n, Eq n, Ord n) => Ord (R1C n) where + compare x@(R1C a b c) y@(R1C e f g) = case (isRank1 x, isRank1 y) of + (True, False) -> LT + (False, True) -> GT + (True, True) -> + -- both are of rank 1 + -- the one with a constant term on the RHS is considered smaller + case (c, g) of + (Left _, Right _) -> LT + (Right _, Left _) -> GT + (Left c0, Left c1) -> compare c0 c1 -- compare the constant terms + (Right v0, Right v1) -> compare v0 v1 -- compare the polynomials + (False, False) -> + -- both are of rank 2 + compare (a, b, c) (e, f, g) + +instance (Show n, Ord n, Eq n, Num n) => Show (R1C n) where + show (R1C aX bX cX) = case (aX, bX, cX) of + (Left 0, _, _) -> "0 = " ++ showVec cX + (_, Left 0, _) -> "0 = " ++ showVec cX + (Left 1, _, _) -> showVec bX ++ " = " ++ showVec cX + (_, Left 1, _) -> showVec aX ++ " = " ++ showVec cX + (_, _, _) -> showVecWithParen aX ++ " * " ++ showVecWithParen bX ++ " = " ++ showVec cX + where + showVec :: (Show n, Ord n, Eq n, Num n) => Either n (Poly n) -> String + showVec (Left c) = show c + showVec (Right xs) = show xs + + -- wrap the string with parenthesis if it has more than 1 term + showVecWithParen :: (Show n, Ord n, Eq n, Num n) => Either n (Poly n) -> String + showVecWithParen (Left c) = showVec (Left c) -- no parenthesis + showVecWithParen (Right xs) = + let termNumber = + IntMap.size (Poly.coeffs xs) + + if Poly.constant xs == 0 + then 0 + else 1 + in if termNumber < 2 + then showVec (Right xs) + else "(" ++ showVec (Right xs) ++ ")" + +-- | See if a R1C is satified by a given assignment +satisfy :: GaloisField a => R1C a -> IntMap a -> Bool +satisfy constraint assignment + | R1C aV bV cV <- constraint = + evaluate aV assignment * evaluate bV assignment == evaluate cV assignment + where + evaluate :: GaloisField a => Either a (Poly a) -> IntMap a -> a + evaluate (Left x) _ = x + evaluate (Right p) w = Poly.evaluate p w + +-- | Free variables in a R1C +freeVars :: R1C n -> IntSet +freeVars (R1C a b c) = freeVarsE a <> freeVarsE b <> freeVarsE c + where + freeVarsE (Left _) = mempty + freeVarsE (Right p) = Poly.vars p + +-- | An R1C is of rank 1 if either side of the multiplication is a constant +isRank1 :: R1C n -> Bool +isRank1 (R1C a b _) = isConstant a || isConstant b + where + isConstant (Left _) = True + isConstant (Right _) = False diff --git a/src/Keelung/Constraint/R1CS.hs b/src/Keelung/Constraint/R1CS.hs new file mode 100644 --- /dev/null +++ b/src/Keelung/Constraint/R1CS.hs @@ -0,0 +1,100 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} + +-- | Constraint system for rank-1 constraints +module Keelung.Constraint.R1CS (R1CS (..), toR1Cs, CNEQ (..)) where + +import Control.DeepSeq (NFData) +import Data.Serialize (Serialize) +import GHC.Generics (Generic) +import Keelung.Constraint.R1C (R1C (..)) +import Keelung.Data.BinRep (BinRep (..)) +import Keelung.Data.Polynomial qualified as Poly +import Keelung.Syntax (Var) +import Keelung.Syntax.Counters + +-------------------------------------------------------------------------------- + +-- | Rank-1 Constraint System +data R1CS n = R1CS + { -- | List of constraints + r1csConstraints :: [R1C n], + -- | Variable bookkeeping + r1csCounters :: Counters, + -- | For restoring CNQZ constraints during R1CS \<=\> ConstraintSystem conversion + r1csCNEQs :: [CNEQ n] + } + deriving (Generic, Eq, NFData, Functor) + +instance Serialize n => Serialize (R1CS n) + +instance (Num n, Eq n, Show n, Ord n) => Show (R1CS n) where + show (R1CS cs counters _) = + "R1CS {\n" + <> prettyConstraints counters cs + <> prettyVariables counters + <> "}" + +-- | Returns 'R1C's from a 'R1CS', including: +-- 1. ordinary constraints +-- 2. Boolean input variable constraints +-- 3. binary representation constraints +toR1Cs :: (Num n, Eq n) => R1CS n -> [R1C n] +toR1Cs (R1CS ordinaryConstraints counters _) = + ordinaryConstraints + <> booleanInputVarConstraints + <> binRepConstraints + where + booleanInputVarConstraints = + let generate (start, end) = + map + ( \var -> + R1C + (Right (Poly.singleVar var)) + (Right (Poly.singleVar var)) + (Right (Poly.singleVar var)) + ) + [start .. end - 1] + in concatMap generate (getBooleanConstraintRanges counters) + + binRepConstraints = + map + ( \(BinRep fVar width bVar) -> + R1C + (Poly.buildEither 0 [(bVar + i, 2 ^ i) | i <- [0 .. width - 1]]) + (Left 1) + (Right (Poly.singleVar fVar)) + ) + (getBinReps counters) + +-------------------------------------------------------------------------------- + +-- | For restoring CNQZ constraints during R1CS \<=\> ConstraintSystem conversion +-- +-- The encoding for constraint @x != y = out@ and some @m@ is: +-- +-- > (x - y) * m = out +-- > (x - y) * (1 - out) = 0 +data CNEQ n + = CNEQ + (Either Var n) + -- ^ @x@: could be a variable or a constant + (Either Var n) + -- ^ @y@: could be a variable or a constant + Var + -- ^ @m@: a constant + deriving + ( Generic, + Eq, + NFData, + Functor + ) + +instance Serialize n => Serialize (CNEQ n) + +instance Show n => Show (CNEQ n) where + show (CNEQ (Left x) (Left y) m) = "Q $" <> show x <> " $" <> show y <> " $" <> show m + show (CNEQ (Left x) (Right y) m) = "Q $" <> show x <> " " <> show y <> " $" <> show m + show (CNEQ (Right x) (Left y) m) = "Q " <> show x <> " $" <> show y <> " $" <> show m + show (CNEQ (Right x) (Right y) m) = "Q " <> show x <> " " <> show y <> " $" <> show m diff --git a/src/Keelung/Data/BinRep.hs b/src/Keelung/Data/BinRep.hs new file mode 100644 --- /dev/null +++ b/src/Keelung/Data/BinRep.hs @@ -0,0 +1,36 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# OPTIONS_HADDOCK hide #-} + +module Keelung.Data.BinRep where + +import Control.DeepSeq (NFData) +import Data.Serialize (Serialize) +import GHC.Generics (Generic) +import Keelung.Syntax (Var) + +-------------------------------------------------------------------------------- + +-- | Relation between a variable and binary representation +-- $var = bit₀ + 2bit₁ + 4bit₂ + 8bit₃ + ... + 2^(n-1)bitₙ +data BinRep = BinRep + { -- | The variable + binRepVar :: Var, + -- | Total number of bits + binRepWidth :: Int, + -- | The starting index of the bits + binRepBitStart :: Var + } + deriving (Eq, Generic, NFData) + +instance Serialize BinRep + +instance Show BinRep where + show (BinRep var 1 index) = "$" <> show var <> " = $" <> show index + show (BinRep var 2 index) = "$" <> show var <> " = $" <> show index <> " + 2$" <> show (index + 1) + show (BinRep var 3 index) = "$" <> show var <> " = $" <> show index <> " + 2$" <> show (index + 1) <> " + 4$" <> show (index + 2) + show (BinRep var 4 index) = "$" <> show var <> " = $" <> show index <> " + 2$" <> show (index + 1) <> " + 4$" <> show (index + 2) <> " + 8$" <> show (index + 3) + show (BinRep var w index) = "$" <> show var <> " = $" <> show index <> " + 2$" <> show (index + 1) <> " + ... + 2^" <> show (w - 1) <> "$" <> show (index + w - 1) + +instance Ord BinRep where + compare (BinRep x _ _) (BinRep y _ _) = compare x y diff --git a/src/Keelung/Data/Bits.hs b/src/Keelung/Data/Bits.hs new file mode 100644 --- /dev/null +++ b/src/Keelung/Data/Bits.hs @@ -0,0 +1,130 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | Like `Data.Bits` but with `Boolean` instead of `Bool` +module Keelung.Data.Bits where + +import Data.Bits qualified as Bits +import Data.Word (Word16, Word32, Word64, Word8) +import GHC.TypeNats (KnownNat) +import Keelung.Syntax + +-- | Bitwise operations on Keelung values +class Bits a where + -- {-# MINIMAL bitWidth #-} + -- bitWidth :: a -> Int + + -- | Bitwise \"and\" + (.&.) :: a -> a -> a + + infixl 8 .&. + + -- | Bitwise \"or\" + (.|.) :: a -> a -> a + + infixl 7 .|. + + -- | Bitwise \"xor\" + (.^.) :: a -> a -> a + + infixl 6 .^. + + -- | Rotates left by @i@ bits if @i@ is positive, or right by @-i@ bits otherwise. + rotate :: a -> Int -> a + + infixl 8 `rotate` + + -- | @'shift' x i@ shifts @x@ left by @i@ bits if @i@ is positive, or right by @-i@ bits otherwise. Vacated bits are filled with 0. + shift :: a -> Int -> a + + infixl 8 `shift` + + -- | Retrieve the i-th bit and return it as Boolean + -- The LSB is the 0-th bit and the MSB is the (n-1)-th bit + -- where n is the bit width + -- You can access the MSB with (-1) because the index is modulo n + (!!!) :: a -> Int -> Boolean + + infixl 9 !!! + + -- | Reverse all the bits in the argument + complement :: a -> a + +-- | Synonym for 'shift' +shiftL :: Bits a => a -> Int -> a +shiftL = shift + +-- | Opposite of 'shiftL' +shiftR :: Bits a => a -> Int -> a +shiftR x i = shiftR x (-i) + +-- | Infix version of 'shiftR'. +(.>>.) :: (Bits a) => a -> Int -> a +(.>>.) = shiftR + +infixl 8 .>>. + +-- | Infix version of 'shiftL'. +(.<<.) :: (Bits a) => a -> Int -> a +(.<<.) = shiftL + +infixl 8 .<<. + +instance Bits Boolean where + (.&.) = And + (.|.) = Or + (.^.) = Xor + rotate x _ = x + shift x 0 = x + shift _ _ = false + x !!! _ = x + complement = Not + +instance KnownNat w => Bits (UInt w) where + (.&.) = AndU + (.|.) = OrU + (.^.) = XorU + rotate expr i = RoLU (widthOf expr) i expr + shift expr i = ShLU (widthOf expr) i expr + (!!!) = BitU + complement = NotU + +-- | Make 'Word8' an instance of 'Bits' +instance Bits Word8 where + (.&.) = (Bits..&.) + (.|.) = (Bits..|.) + (.^.) = Bits.xor + rotate = Bits.rotate + shift = Bits.shift + x !!! i = Boolean (Bits.testBit x i) + complement = Bits.complement + +-- | Make 'Word16' an instance of 'Bits' +instance Bits Word16 where + (.&.) = (Bits..&.) + (.|.) = (Bits..|.) + (.^.) = Bits.xor + rotate = Bits.rotate + shift = Bits.shift + x !!! i = Boolean (Bits.testBit x i) + complement = Bits.complement + +-- | Make 'Word32' an instance of 'Bits' +instance Bits Word32 where + (.&.) = (Bits..&.) + (.|.) = (Bits..|.) + (.^.) = Bits.xor + rotate = Bits.rotate + shift = Bits.shift + x !!! i = Boolean (Bits.testBit x i) + complement = Bits.complement + +-- | Make 'Word64' an instance of 'Bits' +instance Bits Word64 where + (.&.) = (Bits..&.) + (.|.) = (Bits..|.) + (.^.) = Bits.xor + rotate = Bits.rotate + shift = Bits.shift + x !!! i = Boolean (Bits.testBit x i) + complement = Bits.complement diff --git a/src/Keelung/Data/N.hs b/src/Keelung/Data/N.hs new file mode 100644 --- /dev/null +++ b/src/Keelung/Data/N.hs @@ -0,0 +1,73 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# OPTIONS_HADDOCK hide #-} + +module Keelung.Data.N where + +import Control.DeepSeq (NFData) +import Data.Euclidean (Euclidean, Field, GcdDomain) +import Data.Field.Galois (GaloisField (..)) +import Data.Group (Group) +import Data.Semiring (Ring, Semiring) +import Data.Serialize (Serialize) +import GHC.Generics (Generic) +import System.Random (Random) +import Test.QuickCheck (Arbitrary) +import Text.PrettyPrint.Leijen.Text (Pretty) + +-------------------------------------------------------------------------------- + +-- | Data type for displaying field elements nicely +-- Elements in the second half of the field are represented as negative numbers +newtype N a = N {unN :: a} + deriving (Eq, Ord, Generic, NFData) + +instance Serialize a => Serialize (N a) + +deriving instance Bounded n => Bounded (N n) + +deriving instance Group n => Group (N n) + +deriving instance Field n => Field (N n) + +deriving instance Monoid n => Monoid (N n) + +deriving instance Semigroup n => Semigroup (N n) + +deriving instance Euclidean n => Euclidean (N n) + +deriving instance Ring n => Ring (N n) + +deriving instance GcdDomain n => GcdDomain (N n) + +deriving instance Semiring n => Semiring (N n) + +deriving instance Arbitrary n => Arbitrary (N n) + +deriving instance Pretty n => Pretty (N n) + +deriving instance Random n => Random (N n) + +deriving instance (GaloisField n, Integral n) => GaloisField (N n) + +deriving instance Fractional n => Fractional (N n) + +deriving instance Num n => Num (N n) + +deriving instance Enum n => Enum (N n) + +deriving instance Real n => Real (N n) + +instance (GaloisField n, Integral n) => Integral (N n) where + quotRem n m = (N q, N r) + where + (q, r) = quotRem (unN n) (unN m) + toInteger (N x) = + let halfway = fromIntegral (order x `div` 2) + in if x >= halfway + then negate (toInteger (order x) - toInteger x) + else toInteger x + +instance (GaloisField n, Integral n) => Show (N n) where + show = show . toInteger diff --git a/src/Keelung/Data/Polynomial.hs b/src/Keelung/Data/Polynomial.hs new file mode 100644 --- /dev/null +++ b/src/Keelung/Data/Polynomial.hs @@ -0,0 +1,211 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} + +-- | Polynomials over a field for use in constraint systems +module Keelung.Data.Polynomial + ( Poly, + buildEither, + buildEither', + buildMaybe, + singleVar, + bind, + vars, + varSize, + coeffs, + mergeCoeffs, + constant, + view, + renumberVars, + evaluate, + -- + delete, + merge, + negate, + substWithPoly, + substWithVector, + substWithIntMap, + ) +where + +import Control.DeepSeq (NFData) +import Data.IntMap.Strict (IntMap) +import Data.IntMap.Strict qualified as IntMap +import Data.IntSet (IntSet) +import Data.Serialize (Serialize) +import Data.Vector (Vector) +import Data.Vector qualified as Vector +import GHC.Generics (Generic) +import Keelung.Syntax (Var) +import Prelude hiding (negate) +import Prelude qualified + +-- | A Poly is a polynomial of the form @c + c₀x₀ + c₁x₁ ... cₙxₙ = 0@ +-- +-- Invariances: +-- +-- * The coefficients are non-zone +-- * The degree of the polynomial is 1 (there's at least one variable) +data Poly n = Poly !n !(IntMap n) + deriving (Functor, Generic, NFData) + +instance Serialize n => Serialize (Poly n) + +-- | Two polynomials are the same, if they have the same coefficients and variables +-- or one is the negation of the other +instance (Eq n, Num n) => Eq (Poly n) where + (Poly c1 v1) == (Poly c2 v2) = + (c1 == c2 && v1 == v2) + || (c1 == -c2 && v1 == IntMap.map Prelude.negate v2) + +instance (Ord n, Num n) => Ord (Poly n) where + compare (Poly c x) (Poly d y) = + if Poly c x == Poly d y + then EQ + else + let numOfTerms1 = IntMap.size x + if c == 0 then 0 else 1 + numOfTerms2 = IntMap.size y + if d == 0 then 0 else 1 + in compare (numOfTerms1, x, c) (numOfTerms2, y, d) + +instance (Show n, Ord n, Eq n, Num n) => Show (Poly n) where + show (Poly n xs) + | n == 0 = + if firstSign == " + " + then concat restOfTerms + else "- " ++ concat restOfTerms + | otherwise = concat (show n : termStrings) + where + (firstSign, restOfTerms) = case concatMap printTerm $ IntMap.toList xs of + [] -> error "[ panic ] Empty PolyG" + (x' : xs') -> (x', xs') + + termStrings = concatMap printTerm $ IntMap.toList xs + -- return a pair of the sign ("+" or "-") and the string representation + printTerm :: (Show n, Ord n, Eq n, Num n) => (Var, n) -> [String] + printTerm (x, c) + | c == 0 = error "printTerm: coefficient of 0" + | c == 1 = [" + ", "$" ++ show x] + | c == -1 = [" - ", "$" ++ show x] + | c < 0 = [" - ", show (Prelude.negate c) ++ "$" ++ show x] + | otherwise = [" + ", show c ++ "$" ++ show x] + +-- | Create a polynomial from a constant and a list of coefficients. +-- Coefficients of 0 are discarded. +buildEither :: (Num n, Eq n) => n -> [(Var, n)] -> Either n (Poly n) +buildEither c xs = + let xs' = IntMap.filter (0 /=) $ IntMap.fromListWith (+) xs + in if IntMap.null xs' + then Left c + else Right (Poly c xs') + +-- | Create a polynomial from a constant and a list of coefficients. +-- Coefficients of 0 are discarded. +buildEither' :: (Num n, Eq n) => n -> IntMap n -> Either n (Poly n) +buildEither' c xs = + let xs' = IntMap.filter (0 /=) xs + in if IntMap.null xs' + then Left c + else Right (Poly c xs') + +-- | Create a polynomial from a constant and a list of coefficients. +-- Coefficients of 0 are discarded. +buildMaybe :: (Num n, Eq n) => n -> IntMap n -> Maybe (Poly n) +buildMaybe c xs = + let xs' = IntMap.filter (0 /=) xs + in if IntMap.null xs' + then Nothing + else Just (Poly c xs') + +-- | Create a simple binding of a variable to a value +bind :: Num n => Var -> n -> Poly n +bind x n = Poly n (IntMap.singleton x (-1)) + +-- | Create a polynomial from a single variable +singleVar :: Num n => Var -> Poly n +singleVar x = Poly 0 (IntMap.singleton x 1) + +-- | Return the set of variables. +vars :: Poly n -> IntSet +vars = IntMap.keysSet . coeffs + +-- | Number of variables. +varSize :: Poly n -> Int +varSize = IntMap.size . coeffs + +-- | Return the mapping of variables to coefficients. +coeffs :: Poly n -> IntMap n +coeffs (Poly _ xs) = xs + +-- | Merge coefficients of the same variable by adding them up +mergeCoeffs :: (Num n, Eq n) => IntMap n -> IntMap n -> IntMap n +mergeCoeffs xs ys = IntMap.filter (0 /=) $ IntMap.unionWith (+) xs ys + +-- | Return the constant. +constant :: Poly n -> n +constant (Poly c _) = c + +-- | View pattern for 'Poly' +view :: Poly n -> (n, IntMap n) +view (Poly c xs) = (c, xs) + +-- | For renumbering the variables. +renumberVars :: (Var -> Var) -> Poly n -> Poly n +renumberVars f (Poly c xs) = Poly c (IntMap.mapKeys f xs) + +-- | Given an assignment of variables, return the value of the polynomial. +evaluate :: (Num n, Eq n) => Poly n -> IntMap n -> n +evaluate (Poly c xs) assignment = + IntMap.foldlWithKey + (\acc k v -> (v * IntMap.findWithDefault 0 k assignment) + acc) + c + xs + +-- | Delete a variable from the polynomial. +delete :: (Num n, Eq n) => Var -> Poly n -> Maybe (Poly n) +delete x (Poly c xs) = buildMaybe c (IntMap.delete x xs) + +-- | Merge two polynomials. +merge :: (Num n, Eq n) => Poly n -> Poly n -> Maybe (Poly n) +merge (Poly c xs) (Poly d ys) = buildMaybe (c + d) (mergeCoeffs xs ys) + +-- | Negate a polynomial. +negate :: (Num n, Eq n) => Poly n -> Poly n +negate (Poly c xs) = Poly (-c) (fmap Prelude.negate xs) + +-- | Substitute a variable in a polynomial with another polynomial. +substWithPoly :: (Num n, Eq n) => Poly n -> Var -> Poly n -> Maybe (Poly n) +substWithPoly (Poly c xs) var (Poly d ys) = + if IntMap.member var xs + then do + let xs' = ys <> IntMap.delete var xs + buildMaybe (c + d) xs' + else return $ Poly c xs + +-- | Substitute variables in a 'Poly' with a 'Vector' of values. +substWithVector :: (Num n, Eq n) => Poly n -> Vector (Maybe n) -> Either n (Poly n) +substWithVector (Poly c xs) bindings = + let (c', xs') = + IntMap.foldlWithKey' + ( \(is, us) var coeff -> + case bindings Vector.!? var of + Nothing -> (is, IntMap.insert var coeff us) + Just Nothing -> (is, IntMap.insert var coeff us) + Just (Just val) -> ((coeff * val) + is, us) + ) + (c, mempty) + xs + in buildEither' c' xs' + +-- | Substitute variables in a 'Poly' with an 'IntMap' of values. +substWithIntMap :: (Num n, Eq n) => Poly n -> IntMap n -> Either n (Poly n) +substWithIntMap (Poly c xs) bindings = + let (c', xs') = + IntMap.foldlWithKey' + ( \(is, us) var coeff -> + case IntMap.lookup var bindings of + Nothing -> (is, IntMap.insert var coeff us) + Just val -> ((coeff * val) + is, us) + ) + (c, mempty) + xs + in buildEither' c' xs' diff --git a/src/Keelung/Data/Struct.hs b/src/Keelung/Data/Struct.hs new file mode 100644 --- /dev/null +++ b/src/Keelung/Data/Struct.hs @@ -0,0 +1,64 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# OPTIONS_HADDOCK hide #-} + +module Keelung.Data.Struct where + +import Control.DeepSeq (NFData) +import Data.IntMap.Strict (IntMap) +import Data.IntMap.Strict qualified as IntMap +import Data.Serialize (Serialize) +import GHC.Generics (Generic) +import Keelung.Syntax (Width) + +-------------------------------------------------------------------------------- + +-- | Data structure for data associated with different primitive datatypes +data Struct f b u = Struct + { structF :: f, + structB :: b, + structU :: IntMap u + } + deriving (Eq, Show, NFData, Generic) + +instance (Serialize f, Serialize b, Serialize u) => Serialize (Struct f b u) + +instance (Semigroup f, Semigroup b, Semigroup u) => Semigroup (Struct f b u) where + Struct f1 b1 u1 <> Struct f2 b2 u2 = Struct (f1 <> f2) (b1 <> b2) (u1 <> u2) + +instance (Monoid f, Monoid b, Monoid u) => Monoid (Struct f b u) where + mempty = Struct mempty mempty mempty + +updateF :: (x -> y) -> Struct x b u -> Struct y b u +updateF func (Struct f b u) = Struct (func f) b u + +updateB :: (x -> y) -> Struct f x u -> Struct f y u +updateB func (Struct f b u) = Struct f (func b) u + +updateU :: Width -> (x -> x) -> Struct f b x -> Struct f b x +updateU w func (Struct f b u) = Struct f b $ IntMap.adjust func w u + +empty :: (Monoid f, Eq f, Eq b, Monoid b) => Struct f b u -> Bool +empty (Struct f b u) = f == mempty && b == mempty && IntMap.null u + +prettyStruct :: (Show f, Show b, Show u) => String -> Struct (IntMap f) (IntMap b) (IntMap u) -> [String] +prettyStruct suffix (Struct f b u) = + map (\(var, val) -> "$F" <> suffix <> show var ++ " := " <> show val) (IntMap.toList f) + <> map (\(var, val) -> "$B" <> suffix <> show var ++ " := " <> show val) (IntMap.toList b) + <> concatMap (\(width, bindings) -> map (\(var, val) -> "$U" <> suffix <> toSubscript width <> show var ++ " := " <> show val) (IntMap.toList bindings)) (IntMap.toList u) + where + toSubscript :: Int -> String + toSubscript = map go . show + where + go c = case c of + '0' -> '₀' + '1' -> '₁' + '2' -> '₂' + '3' -> '₃' + '4' -> '₄' + '5' -> '₅' + '6' -> '₆' + '7' -> '₇' + '8' -> '₈' + '9' -> '₉' + _ -> c diff --git a/src/Keelung/Error.hs b/src/Keelung/Error.hs new file mode 100644 --- /dev/null +++ b/src/Keelung/Error.hs @@ -0,0 +1,62 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} + +-- | Errors of Keelung +module Keelung.Error where + +import Control.DeepSeq (NFData) +import Data.IntMap (IntMap) +import Data.IntMap qualified as IntMap +import Data.Serialize (Serialize) +import GHC.Generics (Generic) +import Keelung.Heap (Addr) + +-------------------------------------------------------------------------------- + +-- | Errors that can occur when running commands +data Error + = DecodeError String -- Cannot decode the output from the Keelung compiler + | CannotLocateKeelungC -- Cannot locate the Keelung compiler + | CannotLocateProver -- Cannot locate the prover + | CannotLocateVerifier -- Cannot locate the verifier + | CannotReadVersionError -- Cannot read the version of the Keelung compiler + | VersionMismatchError Int Int Int -- The version of the Keelung compiler is not supported + | ElabError ElabError + | CompileError String + deriving (Eq, Generic, NFData) + +instance Show Error where + show (DecodeError err) = "Decode Error: " ++ err + show CannotLocateKeelungC = "Cannot locate the Keelung compiler" + show CannotLocateProver = "Cannot locate the prover" + show CannotLocateVerifier = "Cannot locate the verifier" + show CannotReadVersionError = "Cannot read the version of the Keelung compiler" + show (VersionMismatchError major minor patch) = + "The version of the Keelung compiler is not supported: \n" + ++ " expected range of version: >= v0.9.0 and < v0.10.0, but got v" + ++ show major + ++ "." + ++ show minor + ++ "." + ++ show patch + show (ElabError err) = "Elaboration Error: " ++ show err + show (CompileError err) = "Compile Error: " ++ err + +instance Serialize Error + +-------------------------------------------------------------------------------- + +-- | Errors that can occur during elaboration +data ElabError = IndexOutOfBoundsError Addr Int (IntMap Int) + deriving (Eq, Generic, NFData) + +instance Serialize ElabError + +instance Show ElabError where + show (IndexOutOfBoundsError addr i array) = + "Index " + ++ show i + ++ " out of bounds for array #" + ++ show addr + ++ " of length " + ++ show (IntMap.size array) diff --git a/src/Keelung/Field.hs b/src/Keelung/Field.hs new file mode 100644 --- /dev/null +++ b/src/Keelung/Field.hs @@ -0,0 +1,76 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE RankNTypes #-} + +-- | Field types provided by the compiler +module Keelung.Field + ( B64, + GF181, + BN128, + FieldType (..), + realizeAs, + normalize, + module Keelung.Data.N, + ) +where + +import Control.DeepSeq (NFData) +import Data.Field.Galois (Binary, Prime) +import Data.Serialize (Serialize (..)) +import GHC.Generics (Generic) +import Keelung.Data.N + +-------------------------------------------------------------------------------- + +-- | Binary field of 64 bits +type B64 = Binary 18446744073709551643 + +-- | Prime field of order 1552511030102430251236801561344621993261920897571225601 +type GF181 = Prime 1552511030102430251236801561344621993261920897571225601 + +-- | Barreto-Naehrig curve of 128 bits +type BN128 = Prime 21888242871839275222246405745257275088548364400416034343698204186575808495617 + +instance Serialize B64 where + put = put . toInteger + get = fromInteger <$> get + +instance Serialize GF181 where + put = put . toInteger + get = fromInteger <$> get + +instance Serialize BN128 where + put = put . toInteger + get = fromInteger <$> get + +-------------------------------------------------------------------------------- + +-- | Field types provided by the compiler +data FieldType + = -- | Binary field of 64 bits + B64 + | -- | Prime field of order 181 + GF181 + | -- | Barreto-Naehrig curve of 128 bits + BN128 + deriving + ( Generic, + Eq, + Show, + NFData + ) + +-- | Restore the field type from an 'Integer' +realizeAs :: Num n => FieldType -> Integer -> n +realizeAs B64 n = fromInteger n +realizeAs GF181 n = fromInteger n +realizeAs BN128 n = fromInteger n + +-- | Utility function for normalizing an 'Integer' as some field element +-- the number will be negated if it is on the "upper half" of the field +normalize :: FieldType -> Integer -> Integer +normalize B64 n = toInteger (N (fromIntegral n :: B64)) +normalize GF181 n = toInteger (N (fromIntegral n :: GF181)) +normalize BN128 n = toInteger (N (fromIntegral n :: BN128)) diff --git a/src/Keelung/Heap.hs b/src/Keelung/Heap.hs new file mode 100644 --- /dev/null +++ b/src/Keelung/Heap.hs @@ -0,0 +1,67 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# OPTIONS_HADDOCK hide #-} + +-- | Mutable arrays in Keelung +module Keelung.Heap + ( Addr, + Heap, + ElemType (..), + ArrM (..), + lengthOf, + ) +where + +import Control.DeepSeq (NFData) +import Data.IntMap (IntMap) +import Data.Serialize (Serialize) +import GHC.Generics (Generic) +import Keelung.Syntax (Width) + +-- | A mutable array +data ArrM t = ArrayRef ElemType Int Addr + deriving (Eq) + +-- | Length of a mutable array +lengthOf :: ArrM t -> Int +lengthOf ((ArrayRef _ len _)) = len + +-------------------------------------------------------------------------------- + +-- | An "Address" is also just a integer. +type Addr = Int + +-------------------------------------------------------------------------------- + +-- | A Heap is an mapping of mappings of variables +type Heap = + IntMap + ( ElemType, -- kind of element + IntMap Int -- mapping of index to address of element variables + ) + +-- | Type of elements of a array +data ElemType + = -- | Field elements + ElemF + | -- | Booleans + ElemB + | -- | Unsigned integers + ElemU Width + | -- | Arrays (with type of its elements and its size) + ElemArr ElemType Int + | -- | For empty arrays + EmptyArr + deriving (Show, Eq, Generic, NFData) + +instance Serialize ElemType + +instance Semigroup ElemType where + a <> b = case (a, b) of + (ElemF, ElemF) -> ElemF + (ElemB, ElemB) -> ElemB + (ElemArr a' l, ElemArr b' _) -> ElemArr (a' <> b') l + (ElemArr a' l, EmptyArr) -> ElemArr a' l + (EmptyArr, ElemArr b' l) -> ElemArr b' l + (EmptyArr, EmptyArr) -> EmptyArr + _ -> error "ElemType must be the same" diff --git a/src/Keelung/Monad.hs b/src/Keelung/Monad.hs new file mode 100644 --- /dev/null +++ b/src/Keelung/Monad.hs @@ -0,0 +1,645 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE InstanceSigs #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | Monad and statements for building Keelung programs +module Keelung.Monad + ( -- * Monad + Comp, + + -- * Statements + assert, + performDivMod, + assertDivMod, + + -- * Inputs + Proper (..), + freshVarField, + freshVarBool, + freshVarUInt, + InputAccess (..), + inputField, + inputBool, + inputUInt, + inputList2, + inputList3, + inputVec, + inputVec2, + inputVec3, + + -- * Reuse of expressions + Reusable (..), + + -- * Combinators + mapI, + reduce, + + -- * Mutable Array + ArrM, + Mutable, + toArrayM, + fromArrayM, + freeze, + freeze2, + freeze3, + thaw, + thaw2, + thaw3, + accessM, + accessM2, + accessM3, + updateM, + lengthOf, + + -- * Types + Computation (..), + Elaborated (..), + elaborate, + ) +where + +import Control.Arrow (left) +import Control.Monad.Except +import Control.Monad.State.Strict hiding (get, put) +import Data.Data (Proxy (..)) +import Data.IntMap.Strict (IntMap) +import Data.IntMap.Strict qualified as IntMap +import Data.Traversable (mapAccumL) +import Data.Vector (Vector) +import Data.Vector qualified as Vec +import GHC.TypeNats (KnownNat, natVal) +import Keelung.Data.Struct +import Keelung.Error +import Keelung.Heap +import Keelung.Syntax +import Keelung.Syntax.Counters +import Keelung.Syntax.Encode (encode', runHeapM) +import Keelung.Syntax.Encode.Syntax qualified as Encoding + +-------------------------------------------------------------------------------- + +-- | Data structure for elaboration bookkeeping +data Computation = Computation + { -- Variable bookkeeping + compCounters :: !Counters, + -- Size of allocated heap addresses + compAddrSize :: Int, + -- Heap for arrays + compHeap :: Heap, + -- Bindings to expressions + compExprBindings :: Struct (IntMap Field) (IntMap Boolean) (IntMap Encoding.UInt), + -- Assertions are expressions that are expected to be true + compAssertions :: [Boolean], + -- DivMod relations: dividend = divisor * quotient + remainder + compDivModRelsU :: IntMap (Encoding.UInt, Encoding.UInt, Encoding.UInt, Encoding.UInt) + } + deriving (Eq) + +instance Show Computation where + show (Computation _ addrSize _ eb assertions _divModRelsU) = + "{\n" + <> " Address size: " + <> show addrSize + ++ "\n Bindings to expressions: \n" + ++ show eb + ++ "\n Assertions: \n" + ++ show assertions + ++ "\n\ + \}" + +-------------------------------------------------------------------------------- + +-- | The result of elaborating a computation +data Elaborated t = Elaborated + { -- | The resulting expression + elabExpr :: !t, + -- | The state of computation after elaboration + elabComp :: Computation + } + -- = ElaboratedNum Field Computation + deriving (Eq) + +instance Show t => Show (Elaborated t) where + show (Elaborated expr comp) = + "{\n expression: " + ++ show expr + ++ "\n compuation state: \n" + ++ indent (indent (show comp)) + ++ "\n}" + where + indent :: String -> String + indent = unlines . map (" " <>) . lines + +-------------------------------------------------------------------------------- + +-- | The type of a Keelung program +type Comp = StateT Computation (Except ElabError) + +-- | Elaborates a Keelung program +elaborate :: Comp t -> Either Error (Elaborated t) +elaborate prog = do + (expr, comp) <- left ElabError $ runComp (Computation mempty 0 mempty mempty mempty mempty) prog + return $ Elaborated expr comp + +-- | How to run the 'Comp' monad +runComp :: Computation -> Comp a -> Either ElabError (a, Computation) +runComp comp f = runExcept (runStateT f comp) + +modifyCounter :: (Counters -> Counters) -> Comp () +modifyCounter f = modify (\comp -> comp {compCounters = f (compCounters comp)}) + +-------------------------------------------------------------------------------- +-- Variable & Input Variable +-------------------------------------------------------------------------------- + +-- | Modifier for input variables +-- +-- @since 0.8.4.0 +data InputAccess + = -- | For public input variables, visible to the prover and the verifier + Public + | -- | For private input variables, visible to the prover only + Private + +-- | Allocate a fresh 'Field' variable. +-- +-- @since 0.8.4.0 +freshVarF :: Comp Var +freshVarF = do + counters <- gets compCounters + let index = getCount OfIntermediate OfField counters + modifyCounter $ addCount OfIntermediate OfField 1 + return index + +-- | Allocate a fresh 'Boolean' variable. +-- +-- @since 0.8.4.0 +freshVarB :: Comp Var +freshVarB = do + counters <- gets compCounters + let index = getCount OfIntermediate OfBoolean counters + modifyCounter $ addCount OfIntermediate OfBoolean 1 + return index + +-- | Allocate a fresh 'UInt' variable. +-- +-- @since 0.8.4.0 +freshVarU :: Width -> Comp Var +freshVarU width = do + counters <- gets compCounters + let index = getCount OfIntermediate (OfUInt width) counters + modifyCounter $ addCount OfIntermediate (OfUInt width) 1 + return index + +-- | Allocate a fresh input variable. +freshInputVar :: InputAccess -> VarType -> Int -> Comp Var +freshInputVar acc vt n = do + counters <- gets compCounters + case acc of + Public -> do + let index = getCount OfPublicInput vt counters + modifyCounter $ addCount OfPublicInput vt n + return index + Private -> do + let index = getCount OfPrivateInput vt counters + modifyCounter $ addCount OfPrivateInput vt n + return index + +-------------------------------------------------------------------------------- + +-- | Typeclass for operations on base types +class Proper t where + -- | Request a fresh input variable + -- + -- @since 0.1.0.0 + input :: InputAccess -> Comp t + + -- | Request a fresh variable + -- + -- @since 0.8.4.0 + freshVar :: Comp t + + -- | Request a list of fresh input variables + -- default implementation simply applies `replicateM` on `input` + inputList :: InputAccess -> Int -> Comp [t] + inputList acc size = replicateM size $ input acc + + -- | Conditional clause + -- + -- @since 0.1.0.0 + cond :: Boolean -> t -> t -> t + +instance Proper Field where + input = inputField + + -- \| Specialized implementation for Field + inputList acc size = do + start <- freshInputVar acc OfField size + return $ case acc of + Public -> map VarFI [start .. start + size - 1] + Private -> map VarFP [start .. start + size - 1] + + freshVar = VarF <$> freshVarF + + cond = IfF + +instance Proper Boolean where + input = inputBool + + -- \| Specialized implementation for Boolean + inputList acc size = do + start <- freshInputVar acc OfBoolean size + return $ case acc of + Public -> map VarBI [start .. start + size - 1] + Private -> map VarBP [start .. start + size - 1] + + freshVar = VarB <$> freshVarB + + cond = IfB + +instance KnownNat w => Proper (UInt w) where + input = inputUInt + + -- \| Specialized implementation for UInt + inputList acc size = do + start <- freshInputVar acc (OfUInt width) size + return $ case acc of + Public -> map VarUI [start .. start + size - 1] + Private -> map VarUP [start .. start + size - 1] + where + width = fromIntegral (natVal (Proxy :: Proxy w)) + + freshVar = VarU <$> freshVarU width + where + width = fromIntegral (natVal (Proxy :: Proxy w)) + + cond = IfU + +-- | Requests a fresh 'Field' input variable +inputField :: InputAccess -> Comp Field +inputField Public = VarFI <$> freshInputVar Public OfField 1 +inputField Private = VarFP <$> freshInputVar Private OfField 1 + +-- | Requests a fresh 'Boolean' input variable +inputBool :: InputAccess -> Comp Boolean +inputBool Public = VarBI <$> freshInputVar Public OfBoolean 1 +inputBool Private = VarBP <$> freshInputVar Private OfBoolean 1 + +-- | Requests a fresh 'UInt' input variable of some bit width +inputUInt :: forall w. KnownNat w => InputAccess -> Comp (UInt w) +inputUInt acc = case acc of + Public -> VarUI <$> freshInputVar acc (OfUInt width) 1 + Private -> VarUP <$> freshInputVar acc (OfUInt width) 1 + where + width = fromIntegral (natVal (Proxy :: Proxy w)) + +-- | Requests a fresh 'Field' variable +freshVarField :: Comp Field +freshVarField = freshVar + +-- | Requests a fresh 'Boolean' variable +freshVarBool :: Comp Boolean +freshVarBool = freshVar + +-- | Requests a fresh 'UInt' variable of some bit width +freshVarUInt :: KnownNat w => Comp (UInt w) +freshVarUInt = freshVar + +-------------------------------------------------------------------------------- + +-- | Requests a 2D-array of fresh input variables +inputList2 :: Proper t => InputAccess -> Int -> Int -> Comp [[t]] +inputList2 acc sizeM sizeN = replicateM sizeM (inputList acc sizeN) + +-- | Requests a 3D-array of fresh input variables +inputList3 :: Proper t => InputAccess -> Int -> Int -> Int -> Comp [[[t]]] +inputList3 acc sizeM sizeN sizeO = replicateM sizeM (inputList2 acc sizeN sizeO) + +-------------------------------------------------------------------------------- + +-- | Vector version of 'inputList' +inputVec :: Proper t => InputAccess -> Int -> Comp (Vector t) +inputVec acc size = Vec.fromList <$> inputList acc size + +-- | Vector version of 'inputList2' +inputVec2 :: Proper t => InputAccess -> Int -> Int -> Comp (Vector (Vector t)) +inputVec2 acc sizeM sizeN = Vec.fromList <$> replicateM sizeM (inputVec acc sizeN) + +-- | Vector version of 'inputList3' +inputVec3 :: Proper t => InputAccess -> Int -> Int -> Int -> Comp (Vector (Vector (Vector t))) +inputVec3 acc sizeM sizeN sizeO = Vec.fromList <$> replicateM sizeM (inputVec2 acc sizeN sizeO) + +-------------------------------------------------------------------------------- + +-- | Convert a mutable array to a Haskell list +freeze :: Mutable t => ArrM t -> Comp [t] +freeze = fromArrayM + +-- | Convert a mutable 2D-array to a list of lists +freeze2 :: Mutable t => ArrM (ArrM t) -> Comp [[t]] +freeze2 xs = do + xs' <- fromArrayM xs + mapM freeze xs' + +-- | Convert a mutable 3D-array to a list of lists of lists +freeze3 :: Mutable t => ArrM (ArrM (ArrM t)) -> Comp [[[t]]] +freeze3 xs = do + xs' <- fromArrayM xs + mapM freeze2 xs' + +-- | Convert a Haskell list to a mutable array +thaw :: Mutable t => [t] -> Comp (ArrM t) +thaw = toArrayM + +-- | Convert a list of lists to a mutable 2D-array +thaw2 :: Mutable t => [[t]] -> Comp (ArrM (ArrM t)) +thaw2 xs = mapM thaw xs >>= toArrayM + +-- | Convert a list of lists of lists to a mutable 3D-array +thaw3 :: Mutable t => [[[t]]] -> Comp (ArrM (ArrM (ArrM t))) +thaw3 xs = mapM thaw2 xs >>= toArrayM + +-------------------------------------------------------------------------------- + +-- | Typeclass for retrieving the element of an array +class Mutable t where + -- | Allocates a fresh variable for a value + alloc :: t -> Comp Var + + typeOf :: t -> ElemType + + constructElement :: ElemType -> Addr -> t + +instance Mutable Field where + alloc (VarF var) = return var + alloc val = do + var <- freshVarF + assignF var val + return var + + typeOf _ = ElemF + + constructElement ElemF elemAddr = VarF elemAddr + constructElement _ _ = error "expecting element to be of Num" + +instance Mutable Boolean where + alloc (VarB var) = return var + alloc val = do + var <- freshVarB + assignB var val + return var + + typeOf _ = ElemB + + constructElement ElemB elemAddr = VarB elemAddr + constructElement _ _ = error "expecting element to be of Bool" + +instance KnownNat w => Mutable (UInt w) where + alloc (VarU var) = return var + alloc val = do + let width = widthOf val + var <- freshVarU width + heap <- gets compHeap + let encoded = runHeapM heap (encode' val) + assignU width var encoded + return var + + typeOf val = ElemU (widthOf val) + + constructElement (ElemU _) elemAddr = VarU elemAddr + constructElement _ _ = error "expecting element to be of UInt" + +instance Mutable ref => Mutable (ArrM ref) where + alloc xs@((ArrayRef elemType len _)) = do + elements <- mapM (accessM xs) [0 .. len - 1] + fst <$> allocArray elemType elements + + typeOf ((ArrayRef elemType len _)) = ElemArr elemType len + + constructElement (ElemArr l k) elemAddr = ArrayRef l k elemAddr + constructElement EmptyArr elemAddr = ArrayRef EmptyArr 0 elemAddr + constructElement _ _ = error "expecting element to be array" + +-- | Converts a list of values to an 1D-array +toArrayM :: Mutable t => [t] -> Comp (ArrM t) +toArrayM xs = do + if null xs + then snd <$> allocArray EmptyArr xs + else + let kind = typeOf (head xs) + in snd <$> allocArray kind xs + +-- | Convert an array into a list of expressions +fromArrayM :: Mutable t => ArrM t -> Comp [t] +fromArrayM ((ArrayRef _ _ addr)) = readHeapArray addr + +-- | Access an element from a 1-D array +accessM :: Mutable t => ArrM t -> Int -> Comp t +accessM ((ArrayRef _ _ addr)) i = readHeap (addr, i) + +-- | Access an element from a 2-D array +accessM2 :: Mutable t => ArrM (ArrM t) -> (Int, Int) -> Comp t +accessM2 addr (i, j) = accessM addr i >>= flip accessM j + +-- | Access an element from a 3-D array +accessM3 :: Mutable t => ArrM (ArrM (ArrM t)) -> (Int, Int, Int) -> Comp t +accessM3 addr (i, j, k) = accessM addr i >>= flip accessM j >>= flip accessM k + +-- | Update an entry of an array. +updateM :: Mutable t => ArrM t -> Int -> t -> Comp () +updateM (ArrayRef elemType _ addr) i expr = do + var <- alloc expr + writeHeap addr elemType (i, var) + +-------------------------------------------------------------------------------- + +-- | Internal helper function for allocating an array with values +allocArray :: Mutable t => ElemType -> [t] -> Comp (Addr, ArrM u) +allocArray elemType vals = do + -- allocate a new array for holding the variables of these elements + addr <- gets compAddrSize + modify (\st -> st {compAddrSize = succ addr}) + -- allocate new variables for each element + addresses <- mapM alloc vals + let bindings = IntMap.fromDistinctAscList $ zip [0 ..] addresses + modifyHeap (IntMap.insert addr (elemType, bindings)) + return (addr, ArrayRef elemType (length vals) addr) + +-- | Internal helper function for updating an array entry on the heap +writeHeap :: Addr -> ElemType -> (Int, Var) -> Comp () +writeHeap addr elemType (index, ref) = do + let bindings = IntMap.singleton index ref + modifyHeap (IntMap.insertWith (<>) addr (elemType, bindings)) + +modifyHeap :: (Heap -> Heap) -> Comp () +modifyHeap f = do + heap <- gets compHeap + let heap' = f heap + modify (\st -> st {compHeap = heap'}) + +-- | Internal helper function for accessing an element of an array on the heap +readHeap :: Mutable t => (Addr, Int) -> Comp t +readHeap (addr, i) = do + heap <- gets compHeap + case IntMap.lookup addr heap of + Nothing -> error "readHeap: address not found" + Just (elemType, array) -> case IntMap.lookup i array of + Nothing -> throwError $ IndexOutOfBoundsError addr i array + Just var -> return $ constructElement elemType var + +-- | Internal helper function for accessing an array on the heap +readHeapArray :: Mutable t => Addr -> Comp [t] +readHeapArray addr = do + heap <- gets compHeap + case IntMap.lookup addr heap of + Nothing -> error "readHeap: address not found" + Just (elemType, array) -> return $ map (constructElement elemType) (IntMap.elems array) + +-------------------------------------------------------------------------------- + +-- | An alternative to 'foldM' +reduce :: Foldable m => t -> m a -> (t -> a -> Comp t) -> Comp t +reduce a xs f = foldM f a xs + +-- | Map with index, basically @mapi@ in OCaml. +mapI :: Traversable f => (Int -> a -> b) -> f a -> f b +mapI f = snd . mapAccumL (\i x -> (i + 1, f i x)) 0 + +-------------------------------------------------------------------------------- + +-- | Assert that the given expression evaluates to 'true'. +-- +-- Assertions play a central role in Keelung, as Keelung is all about constraints between variables. +-- +-- /Example/ +-- +-- Consider the following program that takes two inputs and asserts that the second input is the square of the first: +-- +-- @ +-- square :: Comp () +-- square = do +-- x <- input +-- y <- input +-- -- assert that \'y\' is the square of \'x\' +-- assert (y `eq` (x * x)) +-- @ +-- +-- @since 0.1.0.0 +assert :: Boolean -> Comp () +assert expr = modify' $ \st -> st {compAssertions = expr : compAssertions st} + +-------------------------------------------------------------------------------- + +-- | Allow an expression to be referenced and reused in the future +class Reusable t where + reuse :: t -> Comp t + +instance Reusable Boolean where + reuse val = do + var <- freshVarB + assignB var val + return (VarB var) + +instance Reusable Field where + reuse val = do + var <- freshVarF + assignF var val + return (VarF var) + +instance KnownNat w => Reusable (UInt w) where + reuse val = do + var <- freshVarU (widthOf val) + heap <- gets compHeap + let encoded = runHeapM heap (encode' val) + assignU (widthOf val) var encoded + return (VarU var) + +instance (Reusable t, Mutable t) => Reusable (ArrM t) where + reuse = return + +instance (Reusable t, Traversable f) => Reusable (f t) where + reuse = mapM reuse + +assignF :: Var -> Field -> Comp () +assignF var expr = modify' $ \st -> st {compExprBindings = updateF (IntMap.insert var expr) (compExprBindings st)} + +assignB :: Var -> Boolean -> Comp () +assignB var expr = modify' $ \st -> st {compExprBindings = updateB (IntMap.insert var expr) (compExprBindings st)} + +assignU :: Width -> Var -> Encoding.UInt -> Comp () +assignU width var expr = modify' $ \st -> st {compExprBindings = updateU width (IntMap.insert var expr) (compExprBindings st)} + +-------------------------------------------------------------------------------- +-- Asserting DivMod relations +-------------------------------------------------------------------------------- + +-- | Computes the quotient and remainder of two 'UInt' arguments: the dividend and the divisor. +-- +-- Note that because 'performDivMod' is a statement, it can only be executed in the 'Comp' context, as shown in the example below: +-- +-- /Example/ +-- +-- @ +-- program :: Comp (UInt 32) +-- program = do +-- dividend <- input +-- divisor <- input +-- (quotient, remainder) <- performDivMod dividend divisor +-- return quotient +-- @ +-- +-- @since 0.8.3.0 +performDivMod :: + forall w. + KnownNat w => + -- | The dividend + UInt w -> + -- | The devisor + UInt w -> + -- | The quotient and remainder + Comp (UInt w, UInt w) +performDivMod dividend divisor = do + remainder <- freshVarU width + quotient <- freshVarU width + assertDivMod dividend divisor (VarU quotient) (VarU remainder) + return (VarU quotient, VarU remainder) + where + width = fromIntegral (natVal (Proxy :: Proxy w)) + +-- | Instead of computing the quotient and remainder from the dividend and divisor with 'performDivMod', +-- we can enforce a relation between the dividend, divisor, quotient, and remainder in Keelung. +-- +-- For example, we can enforce the dividend to be an even number and obtain the quotient at +-- the same time, as shown below: +-- +-- /Example/ +-- +-- @ +-- assertEven :: UInt 32 -> Comp (UInt 32) +-- assertEven dividend = do +-- quotient <- freshVarUInt +-- assertDivMod dividend 2 quotient 0 +-- return quotient +-- @ +-- +-- @since 0.8.3.0 +assertDivMod :: + forall w. + KnownNat w => + -- | The dividend + UInt w -> + -- | The divisor + UInt w -> + -- | The quotient + UInt w -> + -- | The remainder + UInt w -> + Comp () +assertDivMod dividend divisor quotient remainder = do + heap <- gets compHeap + let encoded = runHeapM heap $ (,,,) <$> encode' dividend <*> encode' divisor <*> encode' quotient <*> encode' remainder + modify (\st -> st {compDivModRelsU = IntMap.insert width encoded (compDivModRelsU st)}) + where + width = fromIntegral (natVal (Proxy :: Proxy w)) diff --git a/src/Keelung/Prof.hs b/src/Keelung/Prof.hs new file mode 100644 --- /dev/null +++ b/src/Keelung/Prof.hs @@ -0,0 +1,17 @@ +{-# OPTIONS_HADDOCK hide #-} + +module Keelung.Prof + ( compileProf, + compileProfWithOpts, + ) +where + +import Keelung +import Keelung.Constraint.R1CS (R1CS) +import Keelung.Error + +compileProf :: Encode t => FieldType -> Comp t -> IO (Either Error (R1CS Integer)) +compileProf = compileProfWithOpts 1 [] [] + +compileProfWithOpts :: Encode t => Int -> [String] -> [String] -> FieldType -> Comp t -> IO (Either Error (R1CS Integer)) +compileProfWithOpts level opts rtsopts = compileWithOpts level opts $ rtsopts <> rtsoptProf diff --git a/src/Keelung/Syntax.hs b/src/Keelung/Syntax.hs new file mode 100644 --- /dev/null +++ b/src/Keelung/Syntax.hs @@ -0,0 +1,300 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | Syntax of the Keelung language +module Keelung.Syntax + ( Field (..), + Boolean (..), + UInt (..), + HasWidth (..), + Cmp (..), + true, + false, + setBit, + Var, + Width, + ) +where + +import Data.Data +import GHC.TypeNats + +-------------------------------------------------------------------------------- + +-- | Field elements. +-- The choice of the underlying field is left to be decided during the compilation. +data Field + = -- | Integral values + Integer Integer + | -- | Rational values + Rational Rational + | -- | Field element variables + VarF Var + | -- | Field element public input variables + VarFI Var + | -- | Field element private input variables + VarFP Var + | -- | Addition + Add Field Field + | -- | Subtraction + Sub Field Field + | -- | Multiplication + Mul Field Field + | -- | Division (without remainders) + Div Field Field + | -- | Conditional that returns a Field element + IfF Boolean Field Field + | -- | Conversion from Booleans to Field elements + BtoF Boolean + +instance Eq Field where + Integer x == Integer y = x == y + Rational x == Rational y = x == y + VarF x == VarF y = x == y + VarFI x == VarFI y = x == y + Add x1 x2 == Add y1 y2 = x1 == y1 && x2 == y2 + Sub x1 x2 == Sub y1 y2 = x1 == y1 && x2 == y2 + Mul x1 x2 == Mul y1 y2 = x1 == y1 && x2 == y2 + Div x1 x2 == Div y1 y2 = x1 == y1 && x2 == y2 + IfF x1 x2 x3 == IfF y1 y2 y3 = x1 == y1 && x2 == y2 && x3 == y3 + BtoF x == BtoF y = x == y + _ == _ = False + +instance Show Field where + showsPrec prec expr = case expr of + Integer n -> showsPrec prec n + Rational n -> showsPrec prec n + VarF ref -> showString "$F" . shows ref + VarFI ref -> showString "$FI" . shows ref + VarFP ref -> showString "$FP" . shows ref + Add x y -> showParen (prec > 6) $ showsPrec 6 x . showString " + " . showsPrec 7 y + Sub x y -> showParen (prec > 6) $ showsPrec 6 x . showString " - " . showsPrec 7 y + Mul x y -> showParen (prec > 7) $ showsPrec 7 x . showString " * " . showsPrec 8 y + Div x y -> showParen (prec > 7) $ showsPrec 7 x . showString " / " . showsPrec 8 y + IfF p x y -> showParen (prec > 1) $ showString "if " . showsPrec 2 p . showString " then " . showsPrec 2 x . showString " else " . showsPrec 2 y + BtoF x -> showString "B→F" . showsPrec prec x + +instance Num Field where + (+) = Add + (-) = Sub + (*) = Mul + abs = id + + -- law of `signum`: abs x * signum x == x + signum = const (Integer 1) + fromInteger = Integer + +instance Fractional Field where + fromRational = Rational + (/) = Div + +-------------------------------------------------------------------------------- + +-- | Unsigned Integers. +-- The bit width is annotated by a type-level natural that is known at compile time. +data UInt (w :: Nat) + = -- | Unsigned integers values + UInt Integer + | -- | Unsigned integer variables + VarU Var + | -- | Unsigned integer public input variables + VarUI Var + | -- | Unsigned integer private input variables + VarUP Var + | -- | Addition + AddU (UInt w) (UInt w) + | -- | Subtraction + SubU (UInt w) (UInt w) + | -- | Multiplication + MulU (UInt w) (UInt w) + | -- | Bitwise conjunction + AndU (UInt w) (UInt w) + | -- | Bitwise disjunction + OrU (UInt w) (UInt w) + | -- | Bitwise exclusive disjunction + XorU (UInt w) (UInt w) + | -- | Bitwise complement + NotU (UInt w) + | -- | Rotate left + RoLU Width Int (UInt w) + | -- | Shift left + ShLU Width Int (UInt w) + | -- | Bit set and return the result + SetU (UInt w) Int Boolean + | -- | Conditional that returns an unsigned integer + IfU Boolean (UInt w) (UInt w) + | -- | Conversion from Booleans to Unsigned integers + BtoU Boolean + deriving (Eq) + +instance KnownNat w => Show (UInt w) where + showsPrec prec expr = case expr of + UInt n -> showsPrec prec n + VarU var -> showString "$U" . showString (toSubscript width) . shows var + VarUI var -> showString "$UI" . showString (toSubscript width) . shows var + VarUP var -> showString "$UP" . showString (toSubscript width) . shows var + AddU x y -> showParen (prec > 6) $ showsPrec 6 x . showString " + " . showsPrec 7 y + SubU x y -> showParen (prec > 6) $ showsPrec 6 x . showString " - " . showsPrec 7 y + MulU x y -> showParen (prec > 7) $ showsPrec 7 x . showString " * " . showsPrec 8 y + AndU x y -> showParen (prec > 5) $ showsPrec 5 x . showString " ∧ " . showsPrec 6 y + OrU x y -> showParen (prec > 4) $ showsPrec 4 x . showString " ∨ " . showsPrec 5 y + XorU x y -> showParen (prec > 3) $ showsPrec 3 x . showString " ⊕ " . showsPrec 4 y + NotU x -> showParen (prec > 8) $ showString "¬ " . showsPrec 9 x + RoLU _ n x -> showParen (prec > 8) $ showString "RoL " . showsPrec 9 n . showString " " . showsPrec 9 x + ShLU _ n x -> showParen (prec > 8) $ showString "ShL " . showsPrec 9 n . showString " " . showsPrec 9 x + SetU x i b -> showParen (prec > 8) $ showsPrec 9 x . showString "[" . showsPrec 9 i . showString "] := " . showsPrec 9 b + IfU p x y -> showParen (prec > 1) $ showString "if " . showsPrec 2 p . showString " then " . showsPrec 2 x . showString " else " . showsPrec 2 y + BtoU x -> showString "B→U " . showsPrec prec x + where + width :: Width + width = widthOf expr + + toSubscript :: Int -> String + toSubscript = map go . show + where + go c = case c of + '0' -> '₀' + '1' -> '₁' + '2' -> '₂' + '3' -> '₃' + '4' -> '₄' + '5' -> '₅' + '6' -> '₆' + '7' -> '₇' + '8' -> '₈' + '9' -> '₉' + _ -> c + +instance KnownNat w => Num (UInt w) where + (+) = AddU + (-) = SubU + (*) = MulU + abs = id + + -- law of `signum`: abs x * signum x == x + signum _ = UInt 1 + + fromInteger n = UInt (fromIntegral n) + +-------------------------------------------------------------------------------- + +-- | Typeclass for deriving the bit width of an expression +class HasWidth a where + -- | Derive the bit width of an expression + widthOf :: a -> Int + +instance KnownNat w => HasWidth (UInt w) where + widthOf _ = fromIntegral $ natVal (Proxy :: Proxy w) + +-------------------------------------------------------------------------------- + +-- | Booleans +data Boolean + = -- | Boolean values + Boolean Bool + | -- | Boolean variables + VarB Var + | -- | Boolean public input variables + VarBI Var + | -- | Boolean private input variables + VarBP Var + | -- | Conjunction + And Boolean Boolean + | -- | Disjunction + Or Boolean Boolean + | -- | Exclusive disjunction + Xor Boolean Boolean + | -- | Complement + Not Boolean + | -- | Equality on Booleans + EqB Boolean Boolean + | -- | Equality on Field elements + EqF Field Field + | -- | Equality on Unsigned integers + forall w. KnownNat w => EqU (UInt w) (UInt w) + | -- | Conditional that returns a Boolean + IfB Boolean Boolean Boolean + | -- | Bit test on Unsigned integers + forall w. KnownNat w => BitU (UInt w) Int + +instance Eq Boolean where + Boolean x == Boolean y = x == y + VarB x == VarB y = x == y + VarBI x == VarBI y = x == y + And x1 x2 == And y1 y2 = x1 == y1 && x2 == y2 + Or x1 x2 == Or y1 y2 = x1 == y1 && x2 == y2 + Xor x1 x2 == Xor y1 y2 = x1 == y1 && x2 == y2 + EqB x1 x2 == EqB y1 y2 = x1 == y1 && x2 == y2 + EqF x1 x2 == EqF y1 y2 = x1 == y1 && x2 == y2 + EqU x1 x2 == EqU y1 y2 = case sameNat x1 x2 of + Just Refl -> case sameNat y1 y2 of + Just Refl -> (x1 == x2) == (y1 == y2) + Nothing -> False + Nothing -> False + IfB x1 x2 x3 == IfB y1 y2 y3 = x1 == y1 && x2 == y2 && x3 == y3 + BitU x1 x2 == BitU y1 y2 = case sameNat x1 y1 of + Just Refl -> x2 == y2 + Nothing -> False + _ == _ = False + +instance Show Boolean where + showsPrec prec expr = case expr of + Boolean b -> showsPrec prec b + VarB ref -> showString "$B" . shows ref + VarBI ref -> showString "$BI" . shows ref + VarBP ref -> showString "$BP" . shows ref + BitU n i -> showsPrec prec n . showString "[" . shows i . showString "]" + EqF x y -> showParen (prec > 5) $ showsPrec 6 x . showString " = " . showsPrec 6 y + And x y -> showParen (prec > 3) $ showsPrec 4 x . showString " ∧ " . showsPrec 3 y + Or x y -> showParen (prec > 2) $ showsPrec 3 x . showString " ∨ " . showsPrec 2 y + Xor x y -> showParen (prec > 4) $ showsPrec 5 x . showString " ⊕ " . showsPrec 4 y + Not x -> showParen (prec > 8) $ showString "¬ " . showsPrec 9 x + EqB x y -> showParen (prec > 5) $ showsPrec 6 x . showString " = " . showsPrec 6 y + EqU x y -> showParen (prec > 5) $ showsPrec 6 x . showString " = " . showsPrec 6 y + IfB p x y -> showParen (prec > 1) $ showString "if " . showsPrec 2 p . showString " then " . showsPrec 2 x . showString " else " . showsPrec 2 y + +-------------------------------------------------------------------------------- + +-- | Smart constructor for 'Boolean True' +true :: Boolean +true = Boolean True + +-- | Smart constructor for 'Boolean False' +false :: Boolean +false = Boolean False + +-- | Set the i-th bit of a Unsigned integer with a Boolean +setBit :: KnownNat w => UInt w -> Int -> Boolean -> UInt w +setBit = SetU + +-------------------------------------------------------------------------------- + +-- | Typeclass for comparing values +class Cmp a where + -- | Equality + eq :: a -> a -> Boolean + + -- | Inequality + neq :: a -> a -> Boolean + +instance Cmp Boolean where + eq = EqB + neq x y = Not (x `eq` y) + +instance Cmp Field where + eq = EqF + neq x y = Not (x `eq` y) + +instance KnownNat w => Cmp (UInt w) where + eq = EqU + neq x y = Not (x `eq` y) + +-------------------------------------------------------------------------------- + +-- | A "Variable" is just a synonym for an 'Int' +type Var = Int + +-- | Bit width +type Width = Int diff --git a/src/Keelung/Syntax/Counters.hs b/src/Keelung/Syntax/Counters.hs new file mode 100644 --- /dev/null +++ b/src/Keelung/Syntax/Counters.hs @@ -0,0 +1,397 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# OPTIONS_HADDOCK hide #-} + +module Keelung.Syntax.Counters + ( Counters (..), + VarType (..), + VarSort (..), + reindex, + getCount, + getCountBySort, + getCountByType, + getTotalCount, + addCount, + -- for constraint generation + getOutputVarRange, + getOutputBinRepRange, + getPublicInputVarRange, + getPrivateInputVarRange, + getBinRepConstraintSize, + getBinReps, + getBooleanConstraintSize, + getBooleanConstraintRanges, + -- for parsing raw inputs + getPublicInputSequence, + getPrivateInputSequence, + -- workaround for variable renumbering + setReducedCount, + -- for pretty printing + prettyConstraints, + prettyVariables, + prettyBooleanConstraints, + prettyBinRepConstraints, + ) +where + +import Control.DeepSeq (NFData) +import Data.IntMap.Strict (IntMap) +import Data.IntMap.Strict qualified as IntMap +import Data.Sequence (Seq) +import Data.Sequence qualified as Seq +import Data.Serialize (Serialize) +import GHC.Generics (Generic) +import Keelung.Data.BinRep (BinRep (..)) +import Keelung.Data.Struct (Struct (..)) + +------------------------------------------------------------------------------ + +type Var = Int + +type Width = Int + +-- | "Types" of variables. +data VarType = OfField | OfBoolean | OfUIntBinRep Width | OfUInt Width + deriving (Generic, NFData, Eq, Show) + +instance Serialize VarType + +-- | "Sorts" of variables. +data VarSort = OfOutput | OfPublicInput | OfPrivateInput | OfIntermediate + +------------------------------------------------------------------------------ + +type SmallCounters = Struct Int Int Int + +binRepSize :: IntMap Int -> Int +binRepSize = IntMap.foldlWithKey' (\acc width size -> acc + width * size) 0 + +uIntSize :: IntMap Int -> Int +uIntSize = sum + +smallCounterSize :: SmallCounters -> Int +smallCounterSize (Struct f b u) = + f + b + binRepSize u + uIntSize u + +-------------------------------------------------------------------------------- + +data Counters = Counters + { countOutput :: !SmallCounters, -- counters for output variables + countPublicInput :: !SmallCounters, -- counters for input variables + countPrivateInput :: !SmallCounters, -- counters for input variables + countIntermediate :: !SmallCounters, -- counters for intermediate variables + countPublicInputSequence :: !(Seq (VarType, Int)), -- Sequence of public input variables + countPrivateInputSequence :: !(Seq (VarType, Int)), -- Sequence of private input variables + countReducedVarHack :: !Int -- HACK, keep track of the number of variables reduced after renumbering + } + deriving (Generic, NFData, Eq, Show) + +instance Serialize Counters + +instance Semigroup Counters where + Counters cOut1 cPubIn1 cPrivIn1 cInt1 cPubInSeq1 cPrivInSeq1 cRed1 <> Counters cOut2 cPubIn2 cPrivIn2 cInt2 cPubInSeq2 cPrivInSeq2 cRed2 = + Counters + (addSmallCounters cOut1 cOut2) + (addSmallCounters cPubIn1 cPubIn2) + (addSmallCounters cPrivIn1 cPrivIn2) + (addSmallCounters cInt1 cInt2) + (cPubInSeq1 <> cPubInSeq2) + (cPrivInSeq1 <> cPrivInSeq2) + (cRed1 + cRed2) + where + addSmallCounters :: SmallCounters -> SmallCounters -> SmallCounters + addSmallCounters (Struct f1 b1 u1) (Struct f2 b2 u2) = + Struct (f1 + f2) (b1 + b2) (IntMap.unionWith (+) u1 u2) + +instance Monoid Counters where + mempty = Counters (Struct 0 0 mempty) (Struct 0 0 mempty) (Struct 0 0 mempty) (Struct 0 0 mempty) mempty mempty 0 + +-------------------------------------------------------------------------------- + +-- | Get the current count for a variable of the given type and sort. +getCount :: VarSort -> VarType -> Counters -> Int +getCount sort typ (Counters o i1 i2 x _ _ _) = + case sort of + OfOutput -> go o + OfPublicInput -> go i1 + OfPrivateInput -> go i2 + OfIntermediate -> go x + where + go :: SmallCounters -> Int + go (Struct f b u) = + case typ of + OfField -> f + OfBoolean -> b + OfUIntBinRep w -> w * IntMap.findWithDefault 0 w u + OfUInt w -> IntMap.findWithDefault 0 w u + +-- | Get the current count for a variable group of the given sort. +getCountBySort :: VarSort -> Counters -> Int +getCountBySort sort (Counters o i1 i2 x _ _ _) = + case sort of + OfOutput -> smallCounterSize o + OfPublicInput -> smallCounterSize i1 + OfPrivateInput -> smallCounterSize i2 + OfIntermediate -> smallCounterSize x + +-- | Get the current count for a variable group of the given type. +getCountByType :: VarType -> Counters -> Int +getCountByType typ (Counters o i1 i2 x _ _ _) = + case typ of + OfField -> structF o + structF i1 + structF i2 + structF x + OfBoolean -> structB o + structB i1 + structB i2 + structB x + OfUIntBinRep _ -> binRepSize (structU o) + binRepSize (structU i1) + binRepSize (structU i2) + binRepSize (structU x) + OfUInt _ -> uIntSize (structU o) + uIntSize (structU i1) + uIntSize (structU i2) + uIntSize (structU x) + +setReducedCount :: Int -> Counters -> Counters +setReducedCount n (Counters o i1 i2 x s1 s2 _) = Counters o i1 i2 x s1 s2 n + +-- | Total count of variables +getTotalCount :: Counters -> Int +getTotalCount (Counters o i1 i2 x _ _ reduced) = + -- 'countReducedVarHack' should only have effect on intermediate variables + (smallCounterSize o + smallCounterSize i1 + smallCounterSize i2) + (0 `max` (smallCounterSize x - reduced)) + +-- | Set the current count for a variable of the given type and sort. +addCount :: VarSort -> VarType -> Int -> Counters -> Counters +addCount sort typ n (Counters o i1 i2 x s1 s2 r) = + case sort of + OfOutput -> Counters (adjustSmallCounters o) i1 i2 x s1 s2 r + OfPublicInput -> Counters o (adjustSmallCounters i1) i2 x (s1 <> newInputSequence) s2 r + OfPrivateInput -> Counters o i1 (adjustSmallCounters i2) x s1 (s2 <> newInputSequence) r + OfIntermediate -> Counters o i1 i2 (adjustSmallCounters x) s1 s2 r + where + adjustSmallCounters :: SmallCounters -> SmallCounters + adjustSmallCounters (Struct f b u) = + case typ of + OfField -> Struct (f + n) b u + OfBoolean -> Struct f (b + n) u + OfUIntBinRep _ -> error "[ panic ] Should use `OfUInt` to adjust the counter instead" + OfUInt w -> Struct f b (IntMap.insertWith (+) w n u) + + oldCount = getCount sort typ (Counters o i1 i2 x s1 s2 r) + + newInputSequence :: Seq (VarType, Int) + newInputSequence = Seq.fromList [(typ, index) | index <- [oldCount .. oldCount + n - 1]] + +-- | For parsing raw inputs +getPublicInputSequence :: Counters -> Seq (VarType, Int) +getPublicInputSequence = countPublicInputSequence + +getPrivateInputSequence :: Counters -> Seq (VarType, Int) +getPrivateInputSequence = countPrivateInputSequence + +-------------------------------------------------------------------------------- + +-- | Re-index variables of different sorts and types +reindex :: Counters -> VarSort -> VarType -> Var -> Var +reindex counters sort typ index = offsetOfSort counters sort + offsetOfType (choose sort counters) typ index + where + choose :: VarSort -> Counters -> SmallCounters + choose OfOutput = countOutput + choose OfPublicInput = countPublicInput + choose OfPrivateInput = countPrivateInput + choose OfIntermediate = countIntermediate + +offsetOfSort :: Counters -> VarSort -> Int +offsetOfSort _ OfOutput = 0 +offsetOfSort counters OfPublicInput = smallCounterSize (countOutput counters) +offsetOfSort counters OfPrivateInput = smallCounterSize (countOutput counters) + smallCounterSize (countPublicInput counters) +offsetOfSort counters OfIntermediate = smallCounterSize (countOutput counters) + smallCounterSize (countPublicInput counters) + smallCounterSize (countPrivateInput counters) + +offsetOfType :: SmallCounters -> VarType -> Int -> Int +offsetOfType _ OfField index = index +offsetOfType (Struct f _ _) OfBoolean index = f + index +offsetOfType (Struct f b u) (OfUIntBinRep width) index = + f + + b + + IntMap.size (IntMap.filterWithKey (\width' _ -> width' < width) u) + + width * index +offsetOfType (Struct f b u) (OfUInt width) index = f + b + binRepSize u + IntMap.size (IntMap.filterWithKey (\width' _ -> width' < width) u) + index + +-------------------------------------------------------------------------------- + +getOutputVarRange :: Counters -> (Int, Int) +getOutputVarRange counters = (offsetOfSort counters OfOutput, offsetOfSort counters OfPublicInput) + +getOutputBinRepRange :: Counters -> (Int, Int) +getOutputBinRepRange counters = + let start = offsetOfSort counters OfOutput + getCount OfOutput OfField counters + getCount OfOutput OfBoolean counters + size = binRepSize (structU (countOutput counters)) + in (start, start + size) + +getPublicInputVarRange :: Counters -> (Int, Int) +getPublicInputVarRange counters = + let inputOffset = offsetOfSort counters OfPublicInput + inputSize = getCountBySort OfPublicInput counters + in (inputOffset, inputOffset + inputSize) + +getPrivateInputVarRange :: Counters -> (Int, Int) +getPrivateInputVarRange counters = + let inputOffset = offsetOfSort counters OfPrivateInput + inputSize = getCountBySort OfPrivateInput counters + in (inputOffset, inputOffset + inputSize) + +-- | Generate one BinRep constraint for each UInt input & output variable +getBinRepConstraintSize :: Counters -> Int +getBinRepConstraintSize (Counters o i1 i2 _ _ _ _) = f o + f i1 + f i2 + where + f (Struct _ _ u) = uIntSize u + +getBinReps :: Counters -> [BinRep] +getBinReps counters@(Counters o i1 i2 x _ _ _) = + fromSmallCounter OfOutput o ++ fromSmallCounter OfPublicInput i1 ++ fromSmallCounter OfPrivateInput i2 ++ fromSmallCounter OfIntermediate x + where + fromSmallCounter :: VarSort -> SmallCounters -> [BinRep] + fromSmallCounter sort (Struct _ _ u) = concatMap (fromPair sort) (IntMap.toList u) + + fromPair :: VarSort -> (Width, Int) -> [BinRep] + fromPair sort (width, count) = + let varOffset = reindex counters sort (OfUInt width) 0 + binRepOffset = reindex counters sort (OfUIntBinRep width) 0 + in [BinRep (varOffset + index) width (binRepOffset + width * index) | index <- [0 .. count - 1]] + +-- | Variables that needed to be constrained to be Boolean +-- 1. Boolean output variables +-- 2. UInt BinReps output variables +-- 3. Boolean input variables +-- 4. UInt BinReps input variables +getBooleanConstraintSize :: Counters -> Int +getBooleanConstraintSize (Counters o i1 i2 _ _ _ _) = f o + f i1 + f i2 + where + f (Struct _ b u) = b + binRepSize u + +-- | Variables that needed to be constrained to be Boolean +-- 1. Boolean output variables +-- 2. UInt BinReps output variables +-- 3. Boolean input variables +-- 4. UInt BinReps input variables +getBooleanConstraintRanges :: Counters -> [(Int, Int)] +getBooleanConstraintRanges counters@(Counters o i1 i2 _ _ _ _) = + mergeSegments [booleanVarRange OfOutput o, booleanVarRange OfPublicInput i1, booleanVarRange OfPrivateInput i2] + where + booleanVarRange :: VarSort -> SmallCounters -> (Int, Int) + booleanVarRange sort (Struct _ b u) = (reindex counters sort OfBoolean 0, reindex counters sort OfBoolean 0 + b + binRepSize u) + + mergeSegments :: [(Int, Int)] -> [(Int, Int)] + mergeSegments [] = [] + mergeSegments [(start, end)] + | end == start = [] + | otherwise = [(start, end)] + mergeSegments ((start, end) : (start', end') : xs) + | end == start = mergeSegments ((start', end') : xs) + | end == start' = mergeSegments ((start, end') : xs) + | otherwise = (start, end) : mergeSegments ((start', end') : xs) + +-------------------------------------------------------------------------------- + +prettyVariables :: Counters -> String +prettyVariables counters@(Counters o i1 i2 _ _ _ _) = + let publicInputOffset = offsetOfSort counters OfPublicInput + privateInputOffset = offsetOfSort counters OfPrivateInput + outputOffset = offsetOfSort counters OfOutput + totalSize = getTotalCount counters + + outputVars = case smallCounterSize o of + 0 -> "" + 1 -> " Output variable : $" <> show outputOffset <> "\n" + n -> " Output variables: $" <> show outputOffset <> " ... $" <> show (outputOffset + n - 1) <> "\n" + publicInputVars = case smallCounterSize i1 of + 0 -> "" + 1 -> " Public Input variable : $" <> show publicInputOffset <> "\n" + n -> " Public Input variables: $" <> show publicInputOffset <> " ... $" <> show (publicInputOffset + n - 1) <> "\n" + privateInputVars = case smallCounterSize i2 of + 0 -> "" + 1 -> " Private Input variable : $" <> show privateInputOffset <> "\n" + n -> " Private Input variables: $" <> show privateInputOffset <> " ... $" <> show (privateInputOffset + n - 1) <> "\n" + in if totalSize == 0 + then "" + else + " Variables (" + <> show totalSize + <> "):\n\n" + <> outputVars + <> publicInputVars + <> privateInputVars + <> "\n" + +prettyConstraints :: Show constraint => Counters -> [constraint] -> String +prettyConstraints counters cs = + showConstraintSummary + <> showOrdinaryConstraints + <> showBooleanConstraints + <> showBinRepConstraints + where + -- sizes of constraint groups + totalBinRepConstraintSize = getBinRepConstraintSize counters + booleanConstraintSize = getBooleanConstraintSize counters + ordinaryConstraintSize = length cs + + -- summary of constraint groups + showConstraintSummary = + " Constriant (" + <> show (ordinaryConstraintSize + booleanConstraintSize + totalBinRepConstraintSize) + <> "): \n" + + -- Ordinary constraints + showOrdinaryConstraints = + if ordinaryConstraintSize == 0 + then "" + else + " Ordinary constriants (" + <> show ordinaryConstraintSize + <> "):\n\n" + <> unlines (map (\x -> " " <> show x) cs) + <> "\n" + + -- Boolean constraints + showBooleanConstraints = + if booleanConstraintSize == 0 + then "" + else + " Boolean constriants (" + <> show booleanConstraintSize + <> "):\n\n" + <> unlines (map (" " <>) (prettyBooleanConstraints counters)) + <> "\n" + + -- BinRep constraints + showBinRepConstraints = + if totalBinRepConstraintSize == 0 + then "" + else + " Binary representation constriants (" + <> show totalBinRepConstraintSize + <> "):\n\n" + <> unlines (map (" " <>) (prettyBinRepConstraints counters)) + <> "\n" + +prettyBooleanConstraints :: Counters -> [String] +prettyBooleanConstraints counters = + concatMap showSegment (getBooleanConstraintRanges counters) + where + showSegment :: (Int, Int) -> [String] + showSegment (start, end) = + case end - start of + 0 -> [] + 1 -> [showBooleanConstraint start] + 2 -> + [ showBooleanConstraint start, + showBooleanConstraint (start + 1) + ] + 3 -> + [ showBooleanConstraint start, + showBooleanConstraint (start + 1), + showBooleanConstraint (start + 2) + ] + _ -> + [ showBooleanConstraint start, + " ...", + showBooleanConstraint (end - 1) + ] + + showBooleanConstraint :: Int -> String + showBooleanConstraint n = "$" <> show n <> " = $" <> show n <> " * $" <> show n + +prettyBinRepConstraints :: Counters -> [String] +prettyBinRepConstraints = map show . getBinReps diff --git a/src/Keelung/Syntax/Encode.hs b/src/Keelung/Syntax/Encode.hs new file mode 100644 --- /dev/null +++ b/src/Keelung/Syntax/Encode.hs @@ -0,0 +1,138 @@ +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | Module for encoding Keelung syntax +module Keelung.Syntax.Encode + ( Encode (..), + runHeapM, + encode', + Expr (..), + ) +where + +import Control.Monad.Reader +import Data.Array.Unboxed qualified as Array +import Data.IntMap qualified as IntMap +import GHC.TypeLits (KnownNat) +import Keelung.Heap +import Keelung.Syntax (widthOf) +import Keelung.Syntax qualified as Syntax +import Keelung.Syntax.Encode.Syntax + +-------------------------------------------------------------------------------- + +-- | MultiParam version of 'Encode' +class Encode' a b where + -- | Encode a Keelung expression + encode' :: a -> HeapM b + +instance Encode' Syntax.Boolean Boolean where + encode' expr = case expr of + Syntax.Boolean b -> return $ ValB b + Syntax.VarB var -> return $ VarB var + Syntax.VarBI var -> return $ VarBI var + Syntax.VarBP var -> return $ VarBP var + Syntax.And x y -> AndB <$> encode' x <*> encode' y + Syntax.Or x y -> OrB <$> encode' x <*> encode' y + Syntax.Xor x y -> XorB <$> encode' x <*> encode' y + Syntax.Not x -> NotB <$> encode' x + Syntax.IfB p x y -> IfB <$> encode' p <*> encode' x <*> encode' y + Syntax.EqB x y -> EqB <$> encode' x <*> encode' y + Syntax.EqF x y -> EqF <$> encode' x <*> encode' y + Syntax.EqU x y -> EqU (widthOf x) <$> encode' x <*> encode' y + Syntax.BitU x i -> BitU (widthOf x) <$> encode' x <*> pure i + +instance Encode' Syntax.Field Field where + encode' expr = case expr of + Syntax.Integer n -> return $ ValF n + Syntax.Rational n -> return $ ValFR n + Syntax.VarF var -> return $ VarF var + Syntax.VarFI var -> return $ VarFI var + Syntax.VarFP var -> return $ VarFP var + Syntax.Add x y -> AddF <$> encode' x <*> encode' y + Syntax.Sub x y -> SubF <$> encode' x <*> encode' y + Syntax.Mul x y -> MulF <$> encode' x <*> encode' y + Syntax.Div x y -> DivF <$> encode' x <*> encode' y + Syntax.IfF p x y -> IfF <$> encode' p <*> encode' x <*> encode' y + Syntax.BtoF b -> BtoF <$> encode' b + +instance KnownNat w => Encode' (Syntax.UInt w) UInt where + encode' expr = case expr of + Syntax.UInt n -> return $ ValU (widthOf expr) n + Syntax.VarU var -> return $ VarU (widthOf expr) var + Syntax.VarUI var -> return $ VarUI (widthOf expr) var + Syntax.VarUP var -> return $ VarUP (widthOf expr) var + Syntax.AddU x y -> AddU (widthOf x) <$> encode' x <*> encode' y + Syntax.SubU x y -> SubU (widthOf x) <$> encode' x <*> encode' y + Syntax.MulU x y -> MulU (widthOf x) <$> encode' x <*> encode' y + Syntax.AndU x y -> AndU (widthOf expr) <$> encode' x <*> encode' y + Syntax.OrU x y -> OrU (widthOf expr) <$> encode' x <*> encode' y + Syntax.XorU x y -> XorU (widthOf expr) <$> encode' x <*> encode' y + Syntax.NotU x -> NotU (widthOf expr) <$> encode' x + Syntax.IfU p x y -> IfU (widthOf expr) <$> encode' p <*> encode' x <*> encode' y + Syntax.RoLU w i x -> RoLU w i <$> encode' x + Syntax.ShLU w i x -> ShLU w i <$> encode' x + Syntax.SetU x i b -> SetU (widthOf expr) <$> encode' x <*> pure i <*> encode' b + Syntax.BtoU n -> BtoU (widthOf expr) <$> encode' n + +-------------------------------------------------------------------------------- + +-- | Typeclass for encoding stuff into something Keelung can understand +class Encode a where + encode :: a -> HeapM Expr + +instance Encode Syntax.Boolean where + encode expr = Boolean <$> encode' expr + +instance Encode Syntax.Field where + encode expr = Field <$> encode' expr + +instance KnownNat w => Encode (Syntax.UInt w) where + encode expr = UInt <$> encode' expr + +instance Encode () where + encode expr = case expr of + () -> return Unit + +instance Encode t => Encode (ArrM t) where + encode expr = case expr of + ArrayRef _ len addr -> readArray addr len + +instance Encode t => Encode [t] where + encode xs = Array . Array.listArray (0, length xs - 1) <$> mapM encode xs + +instance (Encode a, Encode b) => Encode (a, b) where + encode (a, b) = do + a' <- encode a + b' <- encode b + return $ Array $ Array.listArray (0, 1) [a', b'] + +-------------------------------------------------------------------------------- + +-- | Reader Monad for Heap lookups +type HeapM = Reader Heap + +-- | Run a HeapM computation +runHeapM :: Heap -> HeapM a -> a +runHeapM h m = runReader m h + +-- | Read an array from the heap +readArray :: Addr -> Int -> HeapM Expr +readArray addr len = Array <$> mapM (readHeap addr) indices + where + indices :: Array.Array Int Int + indices = Array.listArray (0, pred len) [0 .. pred len] + + readHeap :: Addr -> Int -> HeapM Expr + readHeap addr' i = do + heap <- ask + case IntMap.lookup addr' heap of + Nothing -> error "HeapM: address not found" + Just (elemType, array) -> case IntMap.lookup i array of + Nothing -> error "HeapM: index ouf of bounds" + Just addr'' -> case elemType of + ElemF -> return $ Field $ VarF addr'' + ElemB -> return $ Boolean $ VarB addr'' + ElemU w -> return $ UInt $ VarU w addr'' + ElemArr _ len' -> readArray addr'' len' + EmptyArr -> readArray addr'' 0 diff --git a/src/Keelung/Syntax/Encode/Syntax.hs b/src/Keelung/Syntax/Encode/Syntax.hs new file mode 100644 --- /dev/null +++ b/src/Keelung/Syntax/Encode/Syntax.hs @@ -0,0 +1,277 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} + +-- | Module for encoding Keelung programs +module Keelung.Syntax.Encode.Syntax where + +import Control.DeepSeq (NFData) +import Data.Array.Unboxed (Array) +import Data.Foldable (toList) +import Data.IntMap.Strict (IntMap) +import Data.Serialize (Serialize) +import GHC.Generics (Generic) +import Keelung.Data.Struct +import Keelung.Field (FieldType) +import Keelung.Syntax (Var, Width) +import Keelung.Syntax.Counters + +-------------------------------------------------------------------------------- + +-- | Booleans +data Boolean + = -- | Boolean values + ValB Bool + | -- | Boolean variables + VarB Var + | -- | Boolean public input variables + VarBI Var + | -- | Boolean private input variables + VarBP Var + | -- | Conjunction + AndB Boolean Boolean + | -- | Disjunction + OrB Boolean Boolean + | -- | Exclusive disjunction + XorB Boolean Boolean + | -- | Complement + NotB Boolean + | -- | Conditional that returns a Boolean + IfB Boolean Boolean Boolean + | -- | Equality on Booleans + EqB Boolean Boolean + | -- | Equality on Field elements + EqF Field Field + | -- | Equality on Unsigned integers + EqU Width UInt UInt + | -- | Bit test on Unsigned integers + BitU Width UInt Int + deriving (Generic, Eq, NFData) + +instance Serialize Boolean + +instance Show Boolean where + showsPrec prec expr = case expr of + ValB n -> shows n + VarB var -> showString "$B" . shows var + VarBI var -> showString "$BI" . shows var + VarBP var -> showString "$BP" . shows var + AndB x y -> showParen (prec > 3) $ showsPrec 4 x . showString " ∧ " . showsPrec 3 y + OrB x y -> showParen (prec > 2) $ showsPrec 3 x . showString " ∨ " . showsPrec 2 y + XorB x y -> showParen (prec > 4) $ showsPrec 5 x . showString " ⊕ " . showsPrec 4 y + NotB x -> showParen (prec > 8) $ showString "¬ " . showsPrec 9 x + IfB p x y -> showParen (prec > 1) $ showString "if " . showsPrec 2 p . showString " then " . showsPrec 2 x . showString " else " . showsPrec 2 y + EqB x y -> showParen (prec > 5) $ showsPrec 6 x . showString " = " . showsPrec 6 y + EqF x y -> showParen (prec > 5) $ showsPrec 6 x . showString " = " . showsPrec 6 y + EqU _ x y -> showParen (prec > 5) $ showsPrec 6 x . showString " = " . showsPrec 6 y + BitU _ x i -> showParen (prec > 6) $ showsPrec 7 x . showString " [" . shows i . showString "]" + +-------------------------------------------------------------------------------- + +-- | Field elements +data Field + = -- | Integral values + ValF Integer + | -- | Rational values + ValFR Rational + | -- | Field element variables + VarF Var + | -- | Field element public input variables + VarFI Var + | -- | Field element private input variables + VarFP Var + | -- | Addition + AddF Field Field + | -- | Subtraction + SubF Field Field + | -- | Multiplication + MulF Field Field + | -- | Division (without remainders) + DivF Field Field + | -- | Conditional that returns a Field element + IfF Boolean Field Field + | -- | Conversion from Boolean to Field element + BtoF Boolean + deriving (Generic, Eq, NFData) + +instance Serialize Field + +instance Show Field where + showsPrec prec expr = case expr of + ValF n -> shows n + ValFR n -> shows n + VarF var -> showString "$F" . shows var + VarFI var -> showString "$FI" . shows var + VarFP var -> showString "$FP" . shows var + AddF x y -> showParen (prec > 6) $ showsPrec 6 x . showString " + " . showsPrec 7 y + SubF x y -> showParen (prec > 6) $ showsPrec 6 x . showString " - " . showsPrec 7 y + MulF x y -> showParen (prec > 7) $ showsPrec 7 x . showString " * " . showsPrec 8 y + DivF x y -> showParen (prec > 7) $ showsPrec 7 x . showString " / " . showsPrec 8 y + IfF p x y -> showParen (prec > 1) $ showString "if " . showsPrec 2 p . showString " then " . showsPrec 2 x . showString " else " . showsPrec 2 y + BtoF x -> showString "B→F " . showsPrec prec x + +-------------------------------------------------------------------------------- + +-- | Unsigned Integers +data UInt + = -- | Unsigned integers values + ValU Width Integer + | -- | Unsigned integer variables + VarU Width Var + | -- | Unsigned integer public input variables + VarUI Width Var + | -- | Unsigned integer private input variables + VarUP Width Var + | -- | Addition + AddU Width UInt UInt + | -- | Subtraction + SubU Width UInt UInt + | -- | Multiplication + MulU Width UInt UInt + | -- | Bitwise conjunction + AndU Width UInt UInt + | -- | Bitwise disjunction + OrU Width UInt UInt + | -- | Bitwise exclusive disjunction + XorU Width UInt UInt + | -- | Bitwise complement + NotU Width UInt + | -- | Rotate left + RoLU Width Int UInt + | -- | Shift left + ShLU Width Int UInt + | -- | Set bit and return the result + SetU Width UInt Int Boolean + | -- | Conditional that returns an Unsigned integer + IfU Width Boolean UInt UInt + | -- | Conversion from Boolean to Unsigned integer + BtoU Width Boolean + deriving (Generic, Eq, NFData) + +instance Serialize UInt + +instance Show UInt where + showsPrec prec expr = case expr of + ValU _ n -> shows n + VarU w var -> showString "$U" . showString (toSubscript w) . shows var + VarUI w var -> showString "$UI" . showString (toSubscript w) . shows var + VarUP w var -> showString "$UP" . showString (toSubscript w) . shows var + AddU _ x y -> showParen (prec > 6) $ showsPrec 6 x . showString " + " . showsPrec 7 y + SubU _ x y -> showParen (prec > 6) $ showsPrec 6 x . showString " - " . showsPrec 7 y + MulU _ x y -> showParen (prec > 7) $ showsPrec 7 x . showString " * " . showsPrec 8 y + AndU _ x y -> showParen (prec > 3) $ showsPrec 4 x . showString " ∧ " . showsPrec 3 y + OrU _ x y -> showParen (prec > 2) $ showsPrec 3 x . showString " ∨ " . showsPrec 2 y + XorU _ x y -> showParen (prec > 4) $ showsPrec 5 x . showString " ⊕ " . showsPrec 4 y + NotU _ x -> showParen (prec > 8) $ showString "¬ " . showsPrec prec x + RoLU _ n x -> showParen (prec > 8) $ showString "RoL " . showsPrec 9 n . showString " " . showsPrec 9 x + ShLU _ n x -> showParen (prec > 8) $ showString "ShL " . showsPrec 9 n . showString " " . showsPrec 9 x + SetU _ x i b -> showParen (prec > 8) $ showsPrec 9 x . showString "[" . shows i . showString "] := " . showsPrec 9 b + IfU _ p x y -> showParen (prec > 1) $ showString "if " . showsPrec 2 p . showString " then " . showsPrec 2 x . showString " else " . showsPrec 2 y + BtoU _ x -> showString "B→U " . showsPrec prec x + where + toSubscript :: Int -> String + toSubscript = map go . show + where + go c = case c of + '0' -> '₀' + '1' -> '₁' + '2' -> '₂' + '3' -> '₃' + '4' -> '₄' + '5' -> '₅' + '6' -> '₆' + '7' -> '₇' + '8' -> '₈' + '9' -> '₉' + _ -> c + +-------------------------------------------------------------------------------- + +-- | Encoding of the Keelung syntax +data Expr + = -- | Unit + Unit + | -- | Booleans + Boolean Boolean + | -- | Field element + Field Field + | -- | Unsigned integers + UInt UInt + | -- | Arrays + Array (Array Int Expr) + deriving (Generic, Eq, NFData) + +instance Show Expr where + showsPrec prec expr = case expr of + Unit -> showString "()" + Boolean bool -> showsPrec prec bool + Field num -> showsPrec prec num + UInt uint -> showsPrec prec uint + Array xs -> showList (toList xs) + +instance Serialize Expr + +instance Serialize FieldType + +-------------------------------------------------------------------------------- + +-- | Encoding of a Keelung program after elaboration +data Elaborated = Elaborated + { -- | The resulting 'Expr' + elabExpr :: !Expr, + -- | The state of computation after elaboration + elabComp :: Computation + } + deriving (Generic, NFData) + +instance Show Elaborated where + show (Elaborated expr comp) = + "{\n Expression: \n " + <> showExpr + <> "\n" + <> showExprBindings (compExprBindings comp) + <> showAssertions (compAssertions comp) + <> "}" + where + showExpr = case expr of + Array xs -> prettyList2 4 (toList xs) + _ -> show expr + + showExprBindings eb = + if empty eb + then "" + else + " Bindings of expressions: \n" + <> unlines (map (" " <>) (prettyStruct "" eb)) + <> "\n" + showAssertions assertions = + if null assertions + then "" + else " Assertions: \n" <> unlines (map ((" " <>) . show) assertions) <> "\n" + + prettyList2 :: Show a => Int -> [a] -> String + prettyList2 n list = case list of + [] -> "[]" + [x] -> "[" <> show x <> "]" + (x : xs) -> + unlines $ + map (replicate n ' ' <>) $ + "" : "[ " <> show x : map (\y -> ", " <> show y) xs <> ["]"] + +instance Serialize Elaborated + +-------------------------------------------------------------------------------- + +-- | Data structure for elaboration bookkeeping +data Computation = Computation + { -- Variable bookkeeping + compCounters :: !Counters, + -- Bindings from variables to expressions + compExprBindings :: Struct (IntMap Field) (IntMap Boolean) (IntMap UInt), + -- Assertions are expressions that are expected to be true + compAssertions :: [Expr], + -- DivMod relations: dividend = divisor * quotient + remainder + compDivModRelsU :: IntMap (UInt, UInt, UInt, UInt) + } + deriving (Show, Generic, NFData) + +instance Serialize Computation