diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,567 @@ +# Changelog for Jikka + +## 2021-07-28: v5.0.11.1 + +Uploaded to Hackage + +## 2021-07-21: v5.0.11.0 + +Convex Hull Trick and Segment Trees are implemented. + +Input (`examples/dp_z.py`): + +```python +# https://atcoder.jp/contests/dp/tasks/dp_z +from typing import * + +INF = 10 ** 18 + +def solve(n: int, c: int, h: List[int]) -> int: + assert 2 <= n <= 10 ** 5 + assert 1 <= c <= 10 ** 12 + assert len(h) == n + assert all(1 <= h_i <= 10 ** 6 for h_i in h) + + dp = [INF for _ in range(n)] + dp[0] = 0 + for i in range(1, n): + for j in range(i): + dp[i] = min(dp[i], dp[j] + (h[j] - h[i]) ** 2 + c) + return dp[n - 1] + +def main() -> None: + n, c = map(int, input().split()) + h = list(map(int, input().split())) + assert len(h) == n + ans = solve(n, c, h) + print(ans) + +if __name__ == '__main__': + main() +``` + +Output (): + +```c++ +#include "jikka/base.hpp" +#include "jikka/convex_hull_trick.hpp" +#include +#include +#include +#include +#include +#include +#include +#include +#include +int64_t solve(int64_t n_0, int64_t c_1, std::vector h_2) { + std::vector x3; + x3.push_back(0); + jikka::convex_hull_trick x4_10; + for (int32_t x5 = 0; x5 < n_0 - 1; ++x5) { + x4_10.add_line(h_2[x5], h_2[x5] * h_2[x5] + x3[x5]); + x3.push_back(c_1 + h_2[(x5 + 1)] * h_2[(x5 + 1)] + + x4_10.get_min(h_2[(x5 + 1)] * -2)); + } + return x3[(n_0 - 1)]; +} +int main() { + int64_t n_12 = -1; + int64_t c_13 = -1; + std::cin >> n_12; + std::vector h_14(n_12, -1); + std::cin >> c_13; + for (int32_t i15 = 0; i15 < n_12; ++i15) { + std::cin >> h_14[i15]; + } + auto ans_16 = solve(n_12, c_13, h_14); + std::cout << ans_16 << ' '; + std::cout << '\n' << ' '; +} +``` + +Input (`examples/dp_q.py`): + +```python +# https://atcoder.jp/contests/dp/tasks/dp_q +from typing import * + +def solve(n: int, h: List[int], a: List[int]) -> int: + assert 1 <= n <= 2 * 10 ** 5 + assert len(h) == n + assert all(1 <= h_i <= n for h_i in h) + assert len(a) == n + assert all(1 <= a_i <= 10 ** 9 for a_i in a) + + dp = [0 for _ in range(n)] + for i in range(n): + b = 0 + for j in range(h[i]): + b = max(b, dp[j]) + dp[h[i] - 1] = b + a[i] + return max(dp) + +def main() -> None: + n = int(input()) + h = list(map(int, input().split())) + assert len(h) == n + a = list(map(int, input().split())) + assert len(a) == n + ans = solve(n, h, a) + print(ans) + +if __name__ == '__main__': + main() +``` + +Output (): + +```c++ +#include "jikka/base.hpp" +#include "jikka/segment_tree.hpp" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +int64_t solve(int64_t n_0, std::vector h_1, std::vector a_2) { + std::vector x4(n_0, 0); + atcoder::segtree x5_13( + x4); + for (int32_t x6 = 0; x6 < n_0; ++x6) { + x4[(h_1[x6] - 1)] = std::max(0, x5_13.prod(0, h_1[x6])) + a_2[x6]; + x5_13.set(h_1[x6] - 1, x4[(h_1[x6] - 1)]); + } + int64_t x11 = *std::max_element(x4.begin(), x4.end()); + return x11; +} +int main() { + int64_t n_14 = -1; + std::cin >> n_14; + std::vector h_15(n_14, -1); + std::vector a_16(n_14, -1); + for (int32_t i17 = 0; i17 < n_14; ++i17) { + std::cin >> h_15[i17]; + } + for (int32_t i18 = 0; i18 < n_14; ++i18) { + std::cin >> a_16[i18]; + } + auto ans_19 = solve(n_14, h_15, a_16); + std::cout << ans_19 << ' '; + std::cout << '\n' << ' '; +} +``` + + +## 2021-07-21: v5.0.10.0 + +- The generated C++ code is optimized. +- `list.append` is added in the restricted Python. + +## 2021-07-14: v5.0.9.0 + +- The generated C++ code becomes more natural. +- The restricted Python allows to write `main` function and uses it to analyze input/output format. + +Input: + +```python +# https://atcoder.jp/contests/dp/tasks/dp_a +from typing import * + +def solve(n: int, h: List[int]) -> int: + assert 2 <= n <= 10 ** 5 + assert len(h) == n + assert all(1 <= h_i <= 10 ** 4 for h_i in h) + + dp = [-1 for _ in range(n)] + dp[0] = 0 + dp[1] = abs(h[1] - h[0]) + for i in range(2, n): + dp[i] = min(dp[i - 1] + abs(h[i] - h[i - 1]), dp[i - 2] + abs(h[i] - h[i - 2])) + return dp[n - 1] + +def main() -> None: + n = int(input()) + h = list(map(int, input().split())) + assert len(h) == n + ans = solve(n, h) + print(ans) + +if __name__ == '__main__': + main() +``` + +Output (): + +```c++ +#include "jikka/all.hpp" +#include +#include +#include +#include +#include +#include +#include +#include +int64_t solve(int64_t n_0, std::vector h_1) { + int64_t x6; + if (n_0 + -1 == 0) { + x6 = 0; + } else { + std::vector> x2 = + std::vector>(n_0 + -2 + 1); + x2[0] = jikka::make_array( + 0, std::max(-h_1[0] + h_1[1], h_1[0] + -h_1[1])); + for (int32_t i3 = 0; i3 < int32_t(n_0 + -2); ++i3) { + x2[(i3 + 1)] = jikka::make_array( + x2[i3][1], + std::min( + x2[i3][1] + std::max(-h_1[(i3 + 1)] + h_1[(i3 + 2)], + h_1[(i3 + 1)] + -h_1[(i3 + 2)]), + x2[i3][0] + std::max(-h_1[i3] + h_1[(i3 + 2)], + h_1[i3] + -h_1[(i3 + 2)]))); + } + x6 = x2[(n_0 + -2)][1]; + } + return x6; +} +int main() { + int64_t n_7 = -1; + std::cin >> n_7; + std::vector h_8 = std::vector(n_7, -1); + for (int32_t i9 = 0; i9 < n_7; ++i9) { + std::cin >> h_8[i9]; + } + auto ans_10 = solve(n_7, h_8); + std::cout << ans_10 << ' '; + std::cout << '\n' << ' '; +} +``` + +## 2021-07-11: v5.0.8.0 + +Some optimizers are added. +Now it can use cumulative sums. + +Input: + +```python +# https://judge.yosupo.jp/problem/static_range_sum + +from typing import * + +def solve(n: int, q: int, a: List[int], l: List[int], r: List[int]) -> List[int]: + ans = [-1 for _ in range(q)] + for i in range(q): + ans[i] = sum(a[l[i]:r[i]]) + return ans +``` + +Output (): + +```c++ +std::vector solve(int64_t n_1653, int64_t q_1654, + std::vector a_1655, + std::vector l_1656, + std::vector r_1657) { + std::vector x1658 = jikka::scanl( + [=](int64_t b1659) -> std::function { + return [=](int64_t b1660) -> int64_t { return b1659 + b1660; }; + }, + 0, a_1655); + return jikka::fmap( + [=](int64_t b1661) -> int64_t { + return x1658[(r_1657[b1661] + (-l_1656[b1661] + l_1656[b1661]))] + + -x1658[l_1656[b1661]]; + }, + jikka::range1(q_1654)); +} +``` + +## 2021-07-09: v5.0.7.0 + +Many internal cleanups are done. + +Now our core language is very close to GHC' Core. +It's curried and has a system for rewrite-rules. + + +## 2021-06-29: v5.0.6.0 + +Error reporting and error recovery are improved. + +Input: + +``` python +def solve(n: int) -> bool: + a = n + True # err + b = 2 * n + return b # err +``` + +Output: + +``` console +Type Error (line 2 column 13) (user's mistake?): Jikka.RestrictedPython.Convert.TypeInfer: failed to solve type equations: failed to unify type int and type bool: type int is not type bool +1 |def solve(n: int) -> bool: +2 | a = n + True # err + ^^^^ +3 | b = 2 * n + +Type Error (line 4 column 12) (user's mistake?): Jikka.RestrictedPython.Convert.TypeInfer: failed to solve type equations: failed to unify type bool and type int: type bool is not type int +3 | b = 2 * n +4 | return b # err + ^ +``` + +contributions: + +- @Koki-Yamaguchi fixed build on macOS ([#28](https://github.com/kmyk/Jikka/pull/28)) + + +## 2021-06-25: v5.0.5.0 + +Some optimizations are implemented. +Now it can convert a O(N) Python code for fibonacci to O(log N) C++ code. + +Input, O(N): + +``` python +def f(n: int) -> int: + a = 0 + b = 1 + for _ in range(n): + c = a + b + a = b + b = c + return a + +def solve(n: int) -> int: + return f(n) % 1000000007 +``` + +Output, O(log N): + +``` c++ +#include "jikka/all.hpp" +#include +#include +#include +#include +#include +#include +#include +#include +int64_t solve(int64_t n_317) { + return jikka::modmatap<2, 2>( + jikka::modmatpow<2>(jikka::make_array>( + jikka::make_array(1, 1), + jikka::make_array(1, 0)), + n_317, 1000000007), + jikka::make_array(1, 0), 1000000007)[1]; +} +int main() { + int64_t x318; + std::cin >> x318; + int64_t x319 = solve(x318); + std::cout << x319; + std::cout << '\n'; +} +``` + +## 2021-06-23: v5.0.4.0 + +Now executable C++ code is generated. + +## 2021-06-23: v5.0.3.0 + +Now the conversion from restricted Python to core works. + +## 2021-06-19: v5.0.2.0 + +Most conversions in restricted Python are implemented. + +## 2020-12-03: v5.0.1.0 + +`v5.0.1.0` is the first version of the third prototype. +This version is a more realistic one, which reads a very restricted subset of Python. + +## 2020-04-30: v4.0.1.0 + +`v4.0.1.0` aims to the same thing to `v3.x`, but it's restarted with Haskell. +However, this version didn't reach the usable version. + +## 2019-07-24: v3.1.0 + +`v3.1.0` is the first version which seems to be somewhat usable in practice. + +Input O(N^2): + +``` sml +let given N : [2, 200001) in +let given A : N -> 200001 in + +let f (i : N) = max N (fun j -> if j = i then 0 else A j) in +f +``` + +Output O(N): + +The generated function (+ main function written by hands) gets AC: + +``` c++ +vector solve(int64_t N, const vector & A) { + vector t1(N + 1); + t1[0] = INT64_MIN; + for (int i1 = 0; i1 < N + 1 - 1; ++ i1) { + t1[i1 + 1] = ((0 <= i1) ? max(t1[i1], A[i1]) : INT64_MIN); + } + auto & g1 = t1; + vector t2(N + 1); + t2[0] = INT64_MIN; + for (int i2 = 0; i2 < N + 1 - 1; ++ i2) { + t2[i2 + 1] = ((((N - i2) - 1) < N) ? max(t2[i2], A[((N - i2) - 1)]) : INT64_MIN); + } + auto & g2 = t2; + vector t3(N); + for (int i3 = 0; i3 < N; ++ i3) { + t3[i3] = max(max(max(INT64_MIN, g1[i3]), g2[(((N - (i3 + 1)) - 1) + 1)]), 0); + } + auto & f = t3; + return f; +} +``` + +## 2019-07-19: v3.0.0 + +`v3.0.0` writes C++ function. + +Input O(k n): + +``` sml +# vim: set filetype=sml: +# Jikka v3 + +let K = 100000 in +let given N : Nat in +let given A : Nat -> Nat in + +sum K (fun i -> max N (fun j -> i + 2 * A j)) +``` + +Output O(k + n): + +``` c++ +int64_t solve(int64_t N, const vector & A) { + int64_t K = 100000; + int64_t a2 = 0; + for (int64_t i2 = 0; i2 < K; ++ i2) { + a2 += i2; + } + int64_t a1 = INT64_MIN; + for (int64_t i1 = 0; i1 < N; ++ i1) { + a1 = max(a1, 2 * A[i1]); + } + return a2 + K * a1; +} +``` + +## 2019-07-10: v2 + +> 競技プログラミングの問題の形式的な表現を受けとり、それに対する解法を出力するプログラムです。 + +`v2` is the second prototype. +This version reads a mathematical expression written in ML-like language, and only writes a internal AST. + +Input: + +``` ml +# Jikka v2 +# https://atcoder.jp/contests/code-festival-2015-final-open/tasks/codefestival_2015_final_d + +K = 100000 +input N : Nat +input S : N -> Nat +K1 = K + 1 +input T : N -> K1 +assume forall i. i < N implies S i < T i + +f : N -> K1 -> Nat +f i t = count N (\ j. j < N and j /= i and S j <= t and t < T j) + +output min N \ i. max K1 \ t. f i t +``` + +``` console + +$ dotnet run + +{compiletime = + [(Ident "K1", + AppBExp + (AppBExp + (FreeVarBExp (Ident "+",FunBTy (IntBTy,FunBTy (IntBTy,IntBTy))), + FreeVarBExp (Ident "K",IntBTy)),IntBExp 1), BaseBScm IntBTy); + (Ident "K", IntBExp 100000, BaseBScm IntBTy)]; + input = + [(Ident "T", FunBTy (VarBTy (Ident "N"),VarBTy (Ident "K1"))); + (Ident "S", FunBTy (VarBTy (Ident "N"),VarBTy (Ident "Nat"))); + (Ident "N", VarBTy (Ident "Nat"))]; + runtime = + [...]; + assumptions = + [...]; + output = + (AppBExp + (AppBExp + (FreeVarBExp + (Ident "min",FunBTy (IntBTy,FunBTy (FunBTy (IntBTy,IntBTy),IntBTy))), + FreeVarBExp (Ident "N",IntBTy)), + LamBExp + (IntBTy, + AppBExp + (AppBExp + (FreeVarBExp + (Ident "max", + FunBTy (IntBTy,FunBTy (FunBTy (IntBTy,IntBTy),IntBTy))), + FreeVarBExp (Ident "K1",IntBTy)), + LamBExp + (IntBTy, + AppBExp + (AppBExp + (FreeVarBExp + (Ident "f",FunBTy (IntBTy,FunBTy (IntBTy,IntBTy))), + VarBExp 1),VarBExp 0))))), IntBTy);} +``` + +## 2019-07-02: v1 + +> 数式を入力すると C++ での実装を出力してくれるすごいやつ + +`v1` is the first prototype. +This version reads a mathematical expression written in TeX-like notation, and writes a C++ function. +It is implemented with F#. + +Input: + +``` +\sum _ {i < N} A_i +``` + +Output: + +``` c++ +int64_t solve(const vector & A, int64_t N) { + int64_t t0 = 0; + for (int64_t i = 0; i < N; ++ i) { + t0 += A[i]; + } + return t0; +} +``` diff --git a/Jikka.cabal b/Jikka.cabal new file mode 100644 --- /dev/null +++ b/Jikka.cabal @@ -0,0 +1,259 @@ +cabal-version: 1.12 + +-- This file has been generated from package.yaml by hpack version 0.34.4. +-- +-- see: https://github.com/sol/hpack + +name: Jikka +version: 5.0.11.1 +synopsis: A transpiler from Python to C++ for competitive programming +description: Please see the README on GitHub at +category: Compilers/Interpreters +homepage: https://github.com/kmyk/Jikka#readme +bug-reports: https://github.com/kmyk/Jikka/issues +author: Kimiyuki Onaka +maintainer: kimiyuk95@gmail.com +copyright: 2021 Kimiyuki Onaka +license: Apache +license-file: LICENSE +build-type: Simple +extra-source-files: + README.md + CHANGELOG.md + +source-repository head + type: git + location: https://github.com/kmyk/Jikka + +library + exposed-modules: + Jikka.Common.Alpha + Jikka.Common.Combinatorics + Jikka.Common.Error + Jikka.Common.Format.AutoIndent + Jikka.Common.Format.Color + Jikka.Common.Format.Error + Jikka.Common.Format.Location + Jikka.Common.Format.Show + Jikka.Common.Graph + Jikka.Common.IO + Jikka.Common.IOFormat + Jikka.Common.Location + Jikka.Common.Matrix + Jikka.Common.ModInt + Jikka.Common.Parse.JoinLines + Jikka.Common.Parse.OffsideRule + Jikka.Common.Parse.Read + Jikka.Common.Parse.ShuntingYard + Jikka.Core.Convert + Jikka.Core.Convert.Alpha + Jikka.Core.Convert.ANormal + Jikka.Core.Convert.ArithmeticalExpr + Jikka.Core.Convert.Beta + Jikka.Core.Convert.BubbleLet + Jikka.Core.Convert.CloseAll + Jikka.Core.Convert.CloseMin + Jikka.Core.Convert.CloseSum + Jikka.Core.Convert.ConstantFolding + Jikka.Core.Convert.ConstantPropagation + Jikka.Core.Convert.ConvexHullTrick + Jikka.Core.Convert.CumulativeSum + Jikka.Core.Convert.Eta + Jikka.Core.Convert.MakeScanl + Jikka.Core.Convert.MatrixExponentiation + Jikka.Core.Convert.PropagateMod + Jikka.Core.Convert.RemoveUnusedVars + Jikka.Core.Convert.SegmentTree + Jikka.Core.Convert.ShortCutFusion + Jikka.Core.Convert.SpecializeFoldl + Jikka.Core.Convert.StrengthReduction + Jikka.Core.Convert.TrivialLetElimination + Jikka.Core.Convert.TypeInfer + Jikka.Core.Convert.UnpackTuple + Jikka.Core.Evaluate + Jikka.Core.Format + Jikka.Core.Language.ArithmeticalExpr + Jikka.Core.Language.Beta + Jikka.Core.Language.BuiltinPatterns + Jikka.Core.Language.Expr + Jikka.Core.Language.FreeVars + Jikka.Core.Language.Lint + Jikka.Core.Language.RewriteRules + Jikka.Core.Language.Runtime + Jikka.Core.Language.TypeCheck + Jikka.Core.Language.Util + Jikka.Core.Language.Value + Jikka.CPlusPlus.Convert + Jikka.CPlusPlus.Convert.AddMain + Jikka.CPlusPlus.Convert.FromCore + Jikka.CPlusPlus.Convert.MoveSemantics + Jikka.CPlusPlus.Convert.OptimizeRange + Jikka.CPlusPlus.Convert.UnpackTuples + Jikka.CPlusPlus.Convert.UseInitialization + Jikka.CPlusPlus.Format + Jikka.CPlusPlus.Language.Expr + Jikka.CPlusPlus.Language.Util + Jikka.CPlusPlus.Language.VariableAnalysis + Jikka.Main + Jikka.Main.Subcommand.Convert + Jikka.Main.Subcommand.Debug + Jikka.Main.Subcommand.Execute + Jikka.Main.Target + Jikka.Python.Convert.ToRestrictedPython + Jikka.Python.Language.Expr + Jikka.Python.Language.Util + Jikka.Python.Parse + Jikka.Python.Parse.Alex + Jikka.Python.Parse.Happy + Jikka.Python.Parse.Token + Jikka.RestrictedPython.Convert + Jikka.RestrictedPython.Convert.Alpha + Jikka.RestrictedPython.Convert.DefaultMain + Jikka.RestrictedPython.Convert.ParseMain + Jikka.RestrictedPython.Convert.RemoveUnbalancedIf + Jikka.RestrictedPython.Convert.RemoveUnreachable + Jikka.RestrictedPython.Convert.ResolveBuiltin + Jikka.RestrictedPython.Convert.SplitLoops + Jikka.RestrictedPython.Convert.ToCore + Jikka.RestrictedPython.Convert.TypeInfer + Jikka.RestrictedPython.Convert.UseAppend + Jikka.RestrictedPython.Evaluate + Jikka.RestrictedPython.Format + Jikka.RestrictedPython.Language.Builtin + Jikka.RestrictedPython.Language.Expr + Jikka.RestrictedPython.Language.Lint + Jikka.RestrictedPython.Language.Util + Jikka.RestrictedPython.Language.Value + Jikka.RestrictedPython.Language.VariableAnalysis + Jikka.RestrictedPython.Language.WithoutLoc + other-modules: + Paths_Jikka + hs-source-dirs: + src + ghc-options: -W -optP-Wno-nonportable-include-path + build-tools: + alex + , happy + build-depends: + ansi-terminal ==0.11.* + , array >=0.5.3 && <0.6 + , base ==4.14.* + , containers >=0.6.0 && <0.7 + , deepseq >=1.4.4 && <1.5 + , mtl >=2.2.2 && <2.3 + , text >=1.2.3 && <1.3 + , transformers >=0.5.6 && <0.6 + , vector >=0.12.0 && <0.13 + default-language: Haskell2010 + +executable jikka + main-is: Main.hs + other-modules: + Paths_Jikka + hs-source-dirs: + app + ghc-options: -W -threaded -rtsopts -with-rtsopts=-N -optP-Wno-nonportable-include-path + build-depends: + Jikka + , ansi-terminal ==0.11.* + , array >=0.5.3 && <0.6 + , base ==4.14.* + , containers >=0.6.0 && <0.7 + , deepseq >=1.4.4 && <1.5 + , mtl >=2.2.2 && <2.3 + , text >=1.2.3 && <1.3 + , transformers >=0.5.6 && <0.6 + , vector >=0.12.0 && <0.13 + default-language: Haskell2010 + +test-suite jikka-doctest + type: exitcode-stdio-1.0 + main-is: doctests.hs + other-modules: + Paths_Jikka + hs-source-dirs: + ./ + ghc-options: -W -threaded -rtsopts -with-rtsopts=-N -optP-Wno-nonportable-include-path + build-depends: + Jikka + , ansi-terminal ==0.11.* + , array >=0.5.3 && <0.6 + , base ==4.14.* + , containers >=0.6.0 && <0.7 + , deepseq >=1.4.4 && <1.5 + , doctest + , mtl >=2.2.2 && <2.3 + , text >=1.2.3 && <1.3 + , transformers >=0.5.6 && <0.6 + , vector >=0.12.0 && <0.13 + default-language: Haskell2010 + +test-suite jikka-test + type: exitcode-stdio-1.0 + main-is: Spec.hs + other-modules: + Jikka.Common.MatrixSpec + Jikka.Common.Parse.JoinLinesSpec + Jikka.Common.Parse.OffsideRuleSpec + Jikka.Common.Parse.ShuntingYardSpec + Jikka.Core.Convert.AlphaSpec + Jikka.Core.Convert.ANormalSpec + Jikka.Core.Convert.BetaSpec + Jikka.Core.Convert.CloseSumSpec + Jikka.Core.Convert.ConstantFoldingSpec + Jikka.Core.Convert.ConstantPropagationSpec + Jikka.Core.Convert.EtaSpec + Jikka.Core.Convert.MakeScanlSpec + Jikka.Core.Convert.MatrixExponentiationSpec + Jikka.Core.Convert.PropagateModSpec + Jikka.Core.Convert.RemoveUnusedVarsSpec + Jikka.Core.Convert.ShortCutFusionSpec + Jikka.Core.Convert.SpecializeFoldlSpec + Jikka.Core.Convert.TrivialLetEliminationSpec + Jikka.Core.Convert.TypeInferSpec + Jikka.Core.Convert.UnpackTupleSpec + Jikka.Core.EvaluateSpec + Jikka.Core.FormatSpec + Jikka.Core.Language.ArithmeticalExprSpec + Jikka.Core.Language.BetaSpec + Jikka.CPlusPlus.Convert.FromCoreSpec + Jikka.CPlusPlus.FormatSpec + Jikka.Python.Convert.ToRestrictedPythonSpec + Jikka.Python.Parse.AlexSpec + Jikka.Python.Parse.HappySpec + Jikka.Python.ParseSpec + Jikka.RestrictedPython.Convert.AlphaSpec + Jikka.RestrictedPython.Convert.RemoveUnbalancedIfSpec + Jikka.RestrictedPython.Convert.RemoveUnreachableSpec + Jikka.RestrictedPython.Convert.ResolveBuiltinSpec + Jikka.RestrictedPython.Convert.SplitLoopsSpec + Jikka.RestrictedPython.Convert.ToCoreSpec + Jikka.RestrictedPython.Convert.TypeInferSpec + Jikka.RestrictedPython.EvaluateSpec + Jikka.RestrictedPython.FormatSpec + Jikka.RestrictedPython.Language.BuiltinSpec + Jikka.RestrictedPython.Language.LintSpec + Jikka.RestrictedPython.Language.TypeInferSpec + Jikka.RestrictedPython.Language.UtilSpec + Jikka.RestrictedPython.Language.VariableAnalysisSpec + Paths_Jikka + hs-source-dirs: + test + ghc-options: -W -threaded -rtsopts -with-rtsopts=-N -optP-Wno-nonportable-include-path + build-tool-depends: + hspec-discover:hspec-discover + build-depends: + Jikka + , ansi-terminal ==0.11.* + , array >=0.5.3 && <0.6 + , base ==4.14.* + , containers >=0.6.0 && <0.7 + , deepseq >=1.4.4 && <1.5 + , hlint + , hspec + , mtl >=2.2.2 && <2.3 + , ormolu + , text >=1.2.3 && <1.3 + , transformers >=0.5.6 && <0.6 + , vector >=0.12.0 && <0.13 + default-language: Haskell2010 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 [yyyy] [name of copyright owner] + + 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,179 @@ +# Jikka + +[![test](https://github.com/kmyk/Jikka/actions/workflows/test.yml/badge.svg)](https://github.com/kmyk/Jikka/actions/workflows/test.yml) + +Jikka is an automated solver for problems of competitive programming. + +In competitive programming, there are some problems which are solvable only with "repeating formula transformations", "pasting snippets of famous data structures", etc. +Jikka automatically solves such problems. +Jikka takes such problems as input in the form of a program of a very restricted subset of Python, optimizes the code to reduce the computational complexity, and outputs as an implementation in C++. +/ +競技プログラミングにおいて「ただ式変形をするだけで解ける」「ただデータ構造のライブラリを貼るだけで解ける」問題は実は少なくありません。 +Jikka はそのような問題を自動で解きます。 +そのような問題をとても制限された Python のサブセット言語のコードの形で入力として受け取り、計算量を落とすような最適化を行い、C++ の実装に変換して出力します。 + + +## Usage + +``` console +$ stack run convert PYTHON_FILE +``` + +[Stack](https://www.haskellstack.org/) is required. If you are using Ubuntu, you can install Stack with `$ sudo apt install haskell-stack`. + + +## Documents + +for users: + +- [docs/language.md](https://github.com/kmyk/Jikka/blob/master/docs/language.md) + - [docs/language.ja.md](https://github.com/kmyk/Jikka/blob/master/docs/language.ja.md) (Japanese translation) +- [examples/](https://github.com/kmyk/Jikka/blob/master/examples) +- [CHANGELOG.md](https://github.com/kmyk/Jikka/blob/master/CHANGELOG.md) +- a blog article [競技プログラミングの問題を自動で解きたい - うさぎ小屋](https://kimiyuki.net/blog/2020/12/09/automated-solvers-of-competitive-programming/) (Japanese) + +for developpers: + +- [CONTRIBUTING.md](https://github.com/kmyk/Jikka/blob/master/CONTRIBUTING.md) + - [CONTRIBUTING.ja.md](https://github.com/kmyk/Jikka/blob/master/CONTRIBUTING.ja.md) (Japanese translation) +- [docs/DESIGN.md](https://github.com/kmyk/Jikka/blob/master/docs/DESIGN.md) (Japanese) +- [docs/how-it-works.pdf](https://github.com/kmyk/Jikka/blob/master/docs/how-it-works.pdf) (Japanese) +- [Haddock](https://kmyk.github.io/Jikka/) + + +## Examples + +### `examples/fact.py` (`v5.0.5.0`) + +Input, O(N): + +``` python +def f(n: int) -> int: + a = 0 + b = 1 + for _ in range(n): + c = a + b + a = b + b = c + return a + +def solve(n: int) -> int: + return f(n) % 1000000007 +``` + +Output, O(log N): + +``` c++ +#include "jikka/all.hpp" +#include +#include +#include +#include +#include +#include +#include +#include +int64_t solve(int64_t n_317) { + return jikka::modmatap<2, 2>( + jikka::modmatpow<2>(jikka::make_array>( + jikka::make_array(1, 1), + jikka::make_array(1, 0)), + n_317, 1000000007), + jikka::make_array(1, 0), 1000000007)[1]; +} +int main() { + int64_t x318; + std::cin >> x318; + int64_t x319 = solve(x318); + std::cout << x319; + std::cout << '\n'; +} +``` + +### `examples/static_range_sum.py` (`v5.0.10.0`) + +Input, O(N^2): + +```python +# https://judge.yosupo.jp/problem/static_range_sum + +from typing import * + +def solve(n: int, q: int, a: List[int], l: List[int], r: List[int]) -> List[int]: + ans = [-1 for _ in range(q)] + for i in range(q): + ans[i] = sum(a[l[i]:r[i]]) + return ans + +def main() -> None: + n, q = map(int, input().split()) + a = list(map(int, input().split())) + assert len(a) == n + l = list(range(q)) + r = list(range(q)) + for i in range(q): + l[i], r[i] = map(int, input().split()) + ans = solve(n, q, a, l, r) + for i in range(q): + print(ans[i]) + +if __name__ == '__main__': + main() +``` + +Output, O(N): + +```c++ +#include +#include +#include +#include +#include +#include +#include +#include +std::vector solve(int64_t n_0, int64_t q_1, std::vector a_2, + std::vector l_3, std::vector r_4) { + std::vector x6 = std::vector(a_2.size() + 1); + x6[0] = 0; + for (int32_t i7 = 0; i7 < int32_t(a_2.size()); ++i7) { + x6[(i7 + 1)] = x6[i7] + a_2[i7]; + } + std::vector x5 = x6; + std::vector x10 = std::vector(q_1); + for (int32_t i11 = 0; i11 < int32_t(q_1); ++i11) { + x10[i11] = -x5[l_3[i11]] + x5[r_4[i11]]; + } + return x10; +} +int main() { + int64_t n_13 = -1; + int64_t q_14 = -1; + std::cin >> n_13; + std::vector a_15 = std::vector(n_13, -1); + std::cin >> q_14; + std::vector l_16 = std::vector(q_14, -1); + std::vector r_17 = std::vector(q_14, -1); + for (int32_t i18 = 0; i18 < n_13; ++i18) { + std::cin >> a_15[i18]; + } + for (int32_t i_19 = 0; i_19 < q_14; ++i_19) { + std::cin >> l_16[i_19]; + std::cin >> r_17[i_19]; + } + for (int32_t i_20 = 0; i_20 < q_14; ++i_20) { + } + auto ans_21 = solve(n_13, q_14, a_15, l_16, r_17); + for (int32_t i_22 = 0; i_22 < q_14; ++i_22) { + } + for (int32_t i_23 = 0; i_23 < q_14; ++i_23) { + std::cout << ans_21[i_23] << ' '; + std::cout << '\n' << ' '; + } +} +``` + + +## License + +Appache License 2.0 diff --git a/app/Main.hs b/app/Main.hs new file mode 100644 --- /dev/null +++ b/app/Main.hs @@ -0,0 +1,12 @@ +module Main where + +import qualified Jikka.Main +import System.Environment +import System.Exit + +main :: IO () +main = do + name <- getProgName + args <- getArgs + code <- Jikka.Main.main name args + exitWith code diff --git a/doctests.hs b/doctests.hs new file mode 100644 --- /dev/null +++ b/doctests.hs @@ -0,0 +1,15 @@ +import Test.DocTest + +-- | The modules generated by alex and happy confuse doctest, so we need to avoid some modules which depend on them. +-- TODO: Resolve this issue. +-- +-- > main = doctest ["-isrc", "src/Jikka/Main.hs"] +main :: IO () +main = doctest + [ "src/Jikka/Common/" + , "src/Jikka/Core/" + , "src/Jikka/CPlusPlus/" + , "src/Jikka/Python/Convert/" + , "src/Jikka/Python/Language/" + , "src/Jikka/RestrictedPython/" + ] diff --git a/src/Jikka/CPlusPlus/Convert.hs b/src/Jikka/CPlusPlus/Convert.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/CPlusPlus/Convert.hs @@ -0,0 +1,31 @@ +{-# LANGUAGE FlexibleContexts #-} + +module Jikka.CPlusPlus.Convert + ( run, + ) +where + +import qualified Jikka.CPlusPlus.Convert.AddMain as AddMain +import qualified Jikka.CPlusPlus.Convert.FromCore as FromCore +import qualified Jikka.CPlusPlus.Convert.MoveSemantics as MoveSemantics +import qualified Jikka.CPlusPlus.Convert.OptimizeRange as OptimizeRange +import qualified Jikka.CPlusPlus.Convert.UnpackTuples as UnpackTuples +import qualified Jikka.CPlusPlus.Convert.UseInitialization as UseInitialization +import qualified Jikka.CPlusPlus.Language.Expr as Y +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Common.IOFormat +import qualified Jikka.Core.Language.Expr as X + +run :: (MonadAlpha m, MonadError Error m) => X.Program -> IOFormat -> m Y.Program +run prog format = do + prog <- FromCore.run prog + let go prog = do + prog <- UnpackTuples.run prog + prog <- MoveSemantics.run prog + OptimizeRange.run prog + prog <- go prog + prog <- go prog + prog <- go prog + prog <- AddMain.run prog format + UseInitialization.run prog diff --git a/src/Jikka/CPlusPlus/Convert/AddMain.hs b/src/Jikka/CPlusPlus/Convert/AddMain.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/CPlusPlus/Convert/AddMain.hs @@ -0,0 +1,132 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | +-- Module : Jikka.CPlusPlus.Convert.AddMain +-- Description : adds @main@ function. / @main@ 関数を追加します。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.CPlusPlus.Convert.AddMain + ( run, + ) +where + +import Control.Monad.State.Strict +import qualified Data.Map as M +import qualified Data.Set as S +import Jikka.CPlusPlus.Language.Expr +import Jikka.CPlusPlus.Language.Util +import Jikka.Common.Alpha +import Jikka.Common.Error +import qualified Jikka.Common.IOFormat as F + +lookup' :: (MonadState (M.Map String VarName) m, MonadError Error m) => String -> m VarName +lookup' x = do + y <- gets $ M.lookup x + case y of + Just y -> return y + Nothing -> throwInternalError $ "undefined variable: " ++ x + +runFormatExpr :: (MonadState (M.Map String VarName) m, MonadAlpha m, MonadError Error m) => F.FormatExpr -> m Expr +runFormatExpr = \case + F.Var x -> Var <$> lookup' x + F.Plus e k -> BinOp Add <$> runFormatExpr e <*> pure (Lit (LitInt32 k)) + F.At e i -> at <$> runFormatExpr e <*> (Var <$> lookup' i) + F.Len e -> do + e <- runFormatExpr e + return $ cast TyInt32 (Call MethodSize [e]) + +runMainDeclare :: (MonadState (M.Map String VarName) m, MonadAlpha m, MonadError Error m) => F.IOFormat -> m [(S.Set VarName, Statement)] +runMainDeclare format = go M.empty (F.inputTree format) + where + go sizes = \case + F.Exp e -> do + (x, indices) <- F.unpackSubscriptedVar e + y <- renameVarName LocalNameKind x + modify' $ M.insert x y + let lookupSize i = case M.lookup i sizes of + Just e -> return e + Nothing -> throwInternalError $ "undefined variable" ++ i + sizes' <- mapM lookupSize indices + let deps = S.unions (map freeVars sizes') + let t = foldl (\t _ -> TyVector t) TyInt64 indices + let decl = Declare t y (DeclareCopy (snd (foldr (\size (t, e) -> (TyVector t, vecCtor t [size, e])) (TyInt64, Lit (LitInt64 (-1))) sizes'))) + return [(deps, decl)] + F.Newline -> return [] + F.Seq formats -> concat <$> mapM (go sizes) formats + F.Loop i n body -> do + n <- runFormatExpr n + go (M.insert i n sizes) body + +runMainInput :: (MonadState (M.Map String VarName) m, MonadAlpha m, MonadError Error m) => F.IOFormat -> [(S.Set VarName, Statement)] -> m [Statement] +runMainInput format decls = do + let go initialized = \case + F.Exp e -> do + (x, _) <- F.unpackSubscriptedVar e + y <- lookup' x + e <- runFormatExpr e + let decls' = map snd $ filter (\(deps, _) -> not (deps `S.isSubsetOf` initialized) && deps `S.isSubsetOf` S.insert y initialized) decls + return (cinStatement e : decls', S.insert y initialized) + F.Newline -> return ([], initialized) + F.Seq [] -> return ([], initialized) + F.Seq (format : formats) -> do + (stmts, initialized) <- go initialized format + (stmts', initialized) <- go initialized (F.Seq formats) + return (stmts ++ stmts', initialized) + F.Loop i n body -> do + j <- renameVarName LoopCounterNameKind i + modify' $ M.insert i j + n <- runFormatExpr n + (body, initialized) <- go initialized body + return ([repStatement j n body], initialized) + let decls' = map snd $ filter (\(deps, _) -> S.null deps) decls + stmts <- fst <$> go S.empty (F.inputTree format) + return $ decls' ++ stmts + +runMainSolve :: (MonadState (M.Map String VarName) m, MonadAlpha m, MonadError Error m) => F.IOFormat -> m Statement +runMainSolve format = do + args <- mapM lookup' (F.inputVariables format) + let solve = Call (Function "solve" []) (map Var args) + case F.outputVariables format of + Left x -> do + y <- renameVarName LocalNameKind x + modify' $ M.insert x y + return $ Declare TyAuto y (DeclareCopy solve) + Right xs -> do + ys <- mapM (renameVarName LocalNameKind) xs + modify' $ \env -> foldl (\env (x, y) -> M.insert x y env) env (zip xs ys) + return $ DeclareDestructure ys solve + +runMainOutput :: (MonadState (M.Map String VarName) m, MonadAlpha m, MonadError Error m) => F.IOFormat -> m [Statement] +runMainOutput format = go (F.outputTree format) + where + go = \case + F.Exp e -> do + e <- runFormatExpr e + return [coutStatement e] + F.Newline -> return [coutStatement (Lit (LitChar '\n'))] + F.Seq formats -> concat <$> mapM go formats + F.Loop i n body -> do + j <- renameVarName LoopCounterNameKind i + modify' $ M.insert i j + n <- runFormatExpr n + body <- go body + return [repStatement j n body] + +runMain :: (MonadAlpha m, MonadError Error m) => F.IOFormat -> m ToplevelStatement +runMain format = do + (`evalStateT` M.empty) $ do + decls <- runMainDeclare format + input <- runMainInput format decls + solve <- runMainSolve format + output <- runMainOutput format + return $ FunDef TyInt "main" [] (input ++ [solve] ++ output) + +run :: (MonadAlpha m, MonadError Error m) => Program -> F.IOFormat -> m Program +run prog format = wrapError' "Jikka.CPlusPlus.Convert.AddMain" $ do + main <- runMain format + return $ Program (decls prog ++ [main]) diff --git a/src/Jikka/CPlusPlus/Convert/FromCore.hs b/src/Jikka/CPlusPlus/Convert/FromCore.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/CPlusPlus/Convert/FromCore.hs @@ -0,0 +1,698 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TupleSections #-} + +-- | +-- Module : Jikka.CPlusPlus.Convert.FromCore +-- Description : converts core programs to C++ programs. / core 言語のプログラムを C++ のプログラムに変換します。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- `Jikka.Language.CPlusPlus.FromCore` converts exprs of our core language to exprs of C++. +module Jikka.CPlusPlus.Convert.FromCore + ( run, + ) +where + +import qualified Jikka.CPlusPlus.Language.Expr as Y +import qualified Jikka.CPlusPlus.Language.Util as Y +import Jikka.Common.Alpha +import Jikka.Common.Error +import qualified Jikka.Core.Format as X (formatBuiltinIsolated, formatType) +import qualified Jikka.Core.Language.BuiltinPatterns as X +import qualified Jikka.Core.Language.Expr as X +import qualified Jikka.Core.Language.TypeCheck as X +import qualified Jikka.Core.Language.Util as X + +-------------------------------------------------------------------------------- +-- monad + +renameVarName' :: MonadAlpha m => Y.NameKind -> X.VarName -> m Y.VarName +renameVarName' kind x = Y.renameVarName kind (X.unVarName x) + +type Env = [(X.VarName, X.Type, Y.VarName)] + +typecheckExpr :: MonadError Error m => Env -> X.Expr -> m X.Type +typecheckExpr env = X.typecheckExpr (map (\(x, t, _) -> (x, t)) env) + +lookupVarName :: MonadError Error m => Env -> X.VarName -> m Y.VarName +lookupVarName env x = case lookup x (map (\(x, _, y) -> (x, y)) env) of + Just y -> return y + Nothing -> throwInternalError $ "undefined variable: " ++ X.unVarName x + +-------------------------------------------------------------------------------- +-- run + +runType :: MonadError Error m => X.Type -> m Y.Type +runType = \case + t@X.VarTy {} -> throwInternalError $ "cannot convert type variable: " ++ X.formatType t + X.IntTy -> return Y.TyInt64 + X.BoolTy -> return Y.TyBool + X.ListTy t -> Y.TyVector <$> runType t + X.TupleTy ts -> do + ts <- mapM runType ts + return $ + if Y.shouldBeArray ts + then Y.TyArray (head ts) (fromIntegral (length ts)) + else Y.TyTuple ts + X.FunTy t ret -> Y.TyFunction <$> runType ret <*> mapM runType [t] + X.DataStructureTy ds -> case ds of + X.ConvexHullTrick -> return Y.TyConvexHullTrick + X.SegmentTree semigrp -> return $ Y.TySegmentTree (runSemigroup semigrp) + +runSemigroup :: X.Semigroup' -> Y.Monoid' +runSemigroup = \case + X.SemigroupIntPlus -> Y.MonoidIntPlus + X.SemigroupIntMin -> Y.MonoidIntMin + X.SemigroupIntMax -> Y.MonoidIntMax + +runLiteral :: (MonadAlpha m, MonadError Error m) => Env -> X.Literal -> m Y.Expr +runLiteral env = \case + X.LitBuiltin builtin -> do + (stmts, e) <- runAppBuiltin env builtin [] + case stmts of + [] -> return e + _ -> throwInternalError "now builtin values don't use statements" + X.LitInt n -> return $ Y.Lit (Y.LitInt64 n) + X.LitBool p -> return $ Y.Lit (Y.LitBool p) + X.LitNil t -> do + t <- runType t + return $ Y.vecCtor t [] + X.LitBottom t err -> do + t <- runType t + return $ Y.Call (Y.Function "jikka::error" [t]) [Y.Lit (Y.LitString err)] + +arityOfBuiltin :: X.Builtin -> Int +arityOfBuiltin = \case + X.Min2 _ -> 2 + X.Max2 _ -> 2 + X.Foldl _ _ -> 3 + X.Iterate _ -> 3 + X.At _ -> 2 + X.Min1 _ -> 1 + X.Max1 _ -> 1 + X.Proj _ _ -> 1 + builtin -> length (fst (X.uncurryFunTy (X.builtinToType builtin))) + +runAppBuiltin :: (MonadAlpha m, MonadError Error m) => Env -> X.Builtin -> [X.Expr] -> m ([Y.Statement], Y.Expr) +runAppBuiltin env f args = wrapError' ("converting builtin " ++ X.formatBuiltinIsolated f) $ do + let go0 f = case args of + [] -> return ([], f) + _ -> throwInternalError $ "expected 0 arguments, got " ++ show (length args) + let go1'' :: (MonadAlpha m, MonadError Error m) => (X.Expr -> m ([Y.Statement], Y.Expr)) -> m ([Y.Statement], Y.Expr) + go1'' f = case args of + [e1] -> f e1 + _ -> throwInternalError $ "expected 1 argument, got " ++ show (length args) + let go1' :: (MonadAlpha m, MonadError Error m) => (Y.Expr -> m ([Y.Statement], Y.Expr)) -> m ([Y.Statement], Y.Expr) + go1' f = go1'' $ \e1 -> do + (stmts1, e1) <- runExpr env e1 + (stmts, e) <- f e1 + return (stmts1 ++ stmts, e) + let go1 f = go1' (return . ([],) . f) + let go2'' :: (MonadAlpha m, MonadError Error m) => (X.Expr -> X.Expr -> m ([Y.Statement], Y.Expr)) -> m ([Y.Statement], Y.Expr) + go2'' f = case args of + [e1, e2] -> f e1 e2 + _ -> throwInternalError $ "expected 2 arguments, got " ++ show (length args) + let go2' :: (MonadAlpha m, MonadError Error m) => (Y.Expr -> Y.Expr -> m ([Y.Statement], Y.Expr)) -> m ([Y.Statement], Y.Expr) + go2' f = go2'' $ \e1 e2 -> do + (stmts1, e1) <- runExpr env e1 + (stmts2, e2) <- runExpr env e2 + (stmts, e) <- f e1 e2 + return (stmts1 ++ stmts2 ++ stmts, e) + let go2 f = go2' (((return . ([],)) .) . f) + let go3'' :: (MonadAlpha m, MonadError Error m) => (X.Expr -> X.Expr -> X.Expr -> m ([Y.Statement], Y.Expr)) -> m ([Y.Statement], Y.Expr) + go3'' f = case args of + [e1, e2, e3] -> f e1 e2 e3 + _ -> throwInternalError $ "expected 3 arguments, got " ++ show (length args) + let go3' :: (MonadAlpha m, MonadError Error m) => (Y.Expr -> Y.Expr -> Y.Expr -> m ([Y.Statement], Y.Expr)) -> m ([Y.Statement], Y.Expr) + go3' f = go3'' $ \e1 e2 e3 -> do + (stmts1, e1) <- runExpr env e1 + (stmts2, e2) <- runExpr env e2 + (stmts3, e3) <- runExpr env e3 + (stmts, e) <- f e1 e2 e3 + return (stmts1 ++ stmts2 ++ stmts3 ++ stmts, e) + let go3 f = go3' ((((return . ([],)) .) .) . f) + let goN' :: (MonadAlpha m, MonadError Error m) => ([Y.Expr] -> m Y.Expr) -> m ([Y.Statement], Y.Expr) + goN' f = do + args <- mapM (runExpr env) args + e <- f (map snd args) + return (concatMap fst args, e) + case f of + -- arithmetical functions + X.Negate -> go1 $ \e -> Y.UnOp Y.Negate e + X.Plus -> go2 $ \e1 e2 -> Y.BinOp Y.Add e1 e2 + X.Minus -> go2 $ \e1 e2 -> Y.BinOp Y.Sub e1 e2 + X.Mult -> go2 $ \e1 e2 -> Y.BinOp Y.Mul e1 e2 + X.FloorDiv -> go2 $ \e1 e2 -> Y.Call (Y.Function "jikka::floordiv" []) [e1, e2] + X.FloorMod -> go2 $ \e1 e2 -> Y.Call (Y.Function "jikka::floormod" []) [e1, e2] + X.CeilDiv -> go2 $ \e1 e2 -> Y.Call (Y.Function "jikka::ceildiv" []) [e1, e2] + X.CeilMod -> go2 $ \e1 e2 -> Y.Call (Y.Function "jikka::ceilmod" []) [e1, e2] + X.Pow -> go2 $ \e1 e2 -> Y.Call (Y.Function "jikka::pow" []) [e1, e2] + -- advanced arithmetical functions + X.Abs -> go1 $ \e -> Y.Call (Y.Function "std::abs" []) [e] + X.Gcd -> go2 $ \e1 e2 -> Y.Call (Y.Function "std::gcd" []) [e1, e2] + X.Lcm -> go2 $ \e1 e2 -> Y.Call (Y.Function "std::lcm" []) [e1, e2] + X.Min2 t -> go2' $ \e1 e2 -> do + t <- runType t + return ([], Y.Call (Y.Function "std::min" [t]) [e1, e2]) + X.Max2 t -> go2' $ \e1 e2 -> do + t <- runType t + return ([], Y.Call (Y.Function "std::max" [t]) [e1, e2]) + X.Iterate t -> go3'' $ \n f x -> do + t <- runType t + (stmtsN, n) <- runExpr env n + (stmtsX, x) <- runExpr env x + y <- Y.newFreshName Y.LocalNameKind + i <- Y.newFreshName Y.LoopCounterNameKind + (stmtsF, body, f) <- runExprFunction env f (Y.Var y) + return + ( stmtsN ++ stmtsX + ++ [Y.Declare t y (Y.DeclareCopy x)] + ++ stmtsF + ++ [ Y.repStatement + i + (Y.cast Y.TyInt32 n) + (body ++ [Y.assignSimple y f]) + ], + Y.Var y + ) + -- logical functions + X.Not -> go1 $ \e -> Y.UnOp Y.Not e + X.And -> go2 $ \e1 e2 -> Y.BinOp Y.And e1 e2 + X.Or -> go2 $ \e1 e2 -> Y.BinOp Y.Or e1 e2 + X.Implies -> go2 $ \e1 e2 -> Y.BinOp Y.Or (Y.UnOp Y.Not e1) e2 + X.If t -> go3'' $ \e1 e2 e3 -> do + (stmts1, e1') <- runExpr env e1 + (stmts2, e2') <- runExpr env e2 + (stmts3, e3') <- runExpr env e3 + case (stmts2, stmts3) of + ([], []) + | X.isConstantTimeExpr e2 && X.isConstantTimeExpr e3 -> + return (stmts1, Y.Cond e1' e2' e3') + _ -> do + t <- runType t + phi <- Y.newFreshName Y.LocalNameKind + let assign = Y.Assign . Y.AssignExpr Y.SimpleAssign (Y.LeftVar phi) + return ([Y.Declare t phi Y.DeclareDefault] ++ stmts1 ++ [Y.If e1' (stmts2 ++ [assign e2']) (Just (stmts3 ++ [assign e3']))], Y.Var phi) + -- bitwise functions + X.BitNot -> go1 $ \e -> Y.UnOp Y.BitNot e + X.BitAnd -> go2 $ \e1 e2 -> Y.BinOp Y.BitAnd e1 e2 + X.BitOr -> go2 $ \e1 e2 -> Y.BinOp Y.BitOr e1 e2 + X.BitXor -> go2 $ \e1 e2 -> Y.BinOp Y.BitXor e1 e2 + X.BitLeftShift -> go2 $ \e1 e2 -> Y.BinOp Y.BitLeftShift e1 e2 + X.BitRightShift -> go2 $ \e1 e2 -> Y.BinOp Y.BitRightShift e1 e2 + -- matrix functions + X.MatAp h w -> go2 $ \f x -> Y.Call (Y.Function "jikka::matap" [Y.TyIntValue (fromIntegral h), Y.TyIntValue (fromIntegral w)]) [f, x] + X.MatZero n -> go0 $ Y.Call (Y.Function "jikka::matzero" [Y.TyIntValue (fromIntegral n)]) [] + X.MatOne n -> go0 $ Y.Call (Y.Function "jikka::matone" [Y.TyIntValue (fromIntegral n)]) [] + X.MatAdd h w -> go2 $ \f g -> Y.Call (Y.Function "jikka::matadd" [Y.TyIntValue (fromIntegral h), Y.TyIntValue (fromIntegral w)]) [f, g] + X.MatMul h n w -> go2 $ \f g -> Y.Call (Y.Function "jikka::matmul" [Y.TyIntValue (fromIntegral h), Y.TyIntValue (fromIntegral n), Y.TyIntValue (fromIntegral w)]) [f, g] + X.MatPow n -> go2 $ \f k -> Y.Call (Y.Function "jikka::matpow" [Y.TyIntValue (fromIntegral n)]) [f, k] + X.VecFloorMod n -> go2 $ \x m -> Y.Call (Y.Function "jikka::vecfloormod" [Y.TyIntValue (fromIntegral n)]) [x, m] + X.MatFloorMod h w -> go2 $ \f m -> Y.Call (Y.Function "jikka::matfloormod" [Y.TyIntValue (fromIntegral h), Y.TyIntValue (fromIntegral w)]) [f, m] + -- modular functions + X.ModNegate -> go2 $ \e1 e2 -> Y.Call (Y.Function "jikka::modnegate" []) [e1, e2] + X.ModPlus -> go3 $ \e1 e2 e3 -> Y.Call (Y.Function "jikka::modplus" []) [e1, e2, e3] + X.ModMinus -> go3 $ \e1 e2 e3 -> Y.Call (Y.Function "jikka::modminus" []) [e1, e2, e3] + X.ModMult -> go3 $ \e1 e2 e3 -> Y.Call (Y.Function "jikka::modmult" []) [e1, e2, e3] + X.ModInv -> go2 $ \e1 e2 -> Y.Call (Y.Function "jikka::modinv" []) [e1, e2] + X.ModPow -> go3 $ \e1 e2 e3 -> Y.Call (Y.Function "jikka::modpow" []) [e1, e2, e3] + X.ModMatAp h w -> go3 $ \f x m -> Y.Call (Y.Function "jikka::modmatap" [Y.TyIntValue (fromIntegral h), Y.TyIntValue (fromIntegral w)]) [f, x, m] + X.ModMatAdd h w -> go3 $ \f g m -> Y.Call (Y.Function "jikka::modmatadd" [Y.TyIntValue (fromIntegral h), Y.TyIntValue (fromIntegral w)]) [f, g, m] + X.ModMatMul h n w -> go3 $ \f g m -> Y.Call (Y.Function "jikka::modmatmul" [Y.TyIntValue (fromIntegral h), Y.TyIntValue (fromIntegral n), Y.TyIntValue (fromIntegral w)]) [f, g, m] + X.ModMatPow n -> go3 $ \f k m -> Y.Call (Y.Function "jikka::modmatpow" [Y.TyIntValue (fromIntegral n)]) [f, k, m] + -- list functions + X.Cons t -> go2' $ \x xs -> do + t <- runType t + ys <- Y.newFreshName Y.LocalNameKind + return + ( [ Y.Declare (Y.TyVector t) ys Y.DeclareDefault, + Y.callMethod' (Y.Var ys) "push_back" [x], + Y.callMethod' (Y.Var ys) "insert" [Y.end (Y.Var ys), Y.begin xs, Y.end xs] + ], + Y.Var ys + ) + X.Snoc t -> go2' $ \xs x -> do + t <- runType t + ys <- Y.newFreshName Y.LocalNameKind + return + ( [ Y.Declare (Y.TyVector t) ys (Y.DeclareCopy xs), + Y.callMethod' (Y.Var ys) "push_back" [x] + ], + Y.Var ys + ) + X.Foldl t1 t2 -> go3'' $ \f init xs -> do + (stmtsInit, init) <- runExpr env init + (stmtsXs, xs) <- runExpr env xs + t1 <- runType t1 + t2 <- runType t2 + y <- Y.newFreshName Y.LocalNameKind + x <- Y.newFreshName Y.LocalNameKind + (stmtsF, body, f) <- runExprFunction2 env f (Y.Var y) (Y.Var x) + return + ( stmtsInit ++ stmtsXs + ++ [Y.Declare t2 y (Y.DeclareCopy init)] + ++ stmtsF + ++ [ Y.ForEach + t1 + x + xs + (body ++ [Y.assignSimple y f]) + ], + Y.Var y + ) + X.Scanl _ t2 -> go3'' $ \f init xs -> do + (stmtsInit, init) <- runExpr env init + (stmtsXs, xs) <- runExpr env xs + t2 <- runType t2 + ys <- Y.newFreshName Y.LocalNameKind + i <- Y.newFreshName Y.LoopCounterNameKind + (stmtsF, body, f) <- runExprFunction2 env f (Y.at (Y.Var ys) (Y.Var i)) (Y.at xs (Y.Var i)) + return + ( stmtsInit ++ stmtsXs + ++ [ Y.Declare (Y.TyVector t2) ys (Y.DeclareCopy (Y.vecCtor t2 [Y.incrExpr (Y.size xs)])), + Y.assignAt ys (Y.litInt32 0) init + ] + ++ stmtsF + ++ [ Y.repStatement + i + (Y.cast Y.TyInt32 (Y.size xs)) + (body ++ [Y.assignAt ys (Y.incrExpr (Y.Var i)) f]) + ], + Y.Var ys + ) + X.Build t -> go3'' $ \f xs n -> do + (stmtsInit, xs) <- runExpr env xs + (stmtsXs, n) <- runExpr env n + t <- runType t + ys <- Y.newFreshName Y.LocalNameKind + i <- Y.newFreshName Y.LoopCounterNameKind + (stmtsF, body, f) <- runExprFunction env f (Y.Var ys) + return + ( stmtsInit ++ stmtsXs + ++ [ Y.Declare (Y.TyVector t) ys (Y.DeclareCopy xs) + ] + ++ stmtsF + ++ [ Y.repStatement + i + (Y.cast Y.TyInt32 n) + (body ++ [Y.callMethod' (Y.Var ys) "push_back" [f]]) + ], + Y.Var ys + ) + X.Len _ -> go1 $ \e -> Y.cast Y.TyInt64 (Y.size e) + X.Map _ t2 -> go2'' $ \f xs -> do + ys <- Y.newFreshName Y.LocalNameKind + t2 <- runType t2 + stmts <- case (f, xs) of + (X.Lam _ _ (X.Lit lit), X.Range1' n) -> do + (stmtsN, n) <- runExpr env n + lit <- runLiteral env lit + return $ + stmtsN + ++ [Y.Declare (Y.TyVector t2) ys (Y.DeclareCopy (Y.vecCtor t2 [n, lit]))] + _ -> do + (stmtsXs, xs) <- runExpr env xs + i <- Y.newFreshName Y.LoopCounterNameKind + (stmtsF, body, f) <- runExprFunction env f (Y.at xs (Y.Var i)) + return $ + stmtsXs + ++ [Y.Declare (Y.TyVector t2) ys (Y.DeclareCopy (Y.vecCtor t2 [Y.size xs]))] + ++ stmtsF + ++ [ Y.repStatement + i + (Y.cast Y.TyInt32 (Y.size xs)) + (body ++ [Y.assignAt ys (Y.Var i) f]) + ] + return (stmts, Y.Var ys) + X.Filter t -> go2'' $ \f xs -> do + (stmtsXs, xs) <- runExpr env xs + t <- runType t + ys <- Y.newFreshName Y.LocalNameKind + x <- Y.newFreshName Y.LocalNameKind + (stmtsF, body, f) <- runExprFunction env f (Y.Var x) + return + ( stmtsXs + ++ [Y.Declare (Y.TyVector t) ys Y.DeclareDefault] + ++ stmtsF + ++ [ Y.ForEach + t + x + xs + ( body + ++ [ Y.If + f + [Y.callMethod' (Y.Var ys) "push_back" [Y.Var x]] + Nothing + ] + ) + ], + Y.Var ys + ) + X.At _ -> go2 $ \e1 e2 -> Y.at e1 e2 + X.SetAt t -> go3' $ \xs i x -> do + t <- runType t + ys <- Y.newFreshName Y.LocalNameKind + return + ( [ Y.Declare (Y.TyVector t) ys (Y.DeclareCopy xs), + Y.assignAt ys i x + ], + Y.Var ys + ) + X.Elem _ -> go2' $ \xs x -> do + y <- Y.newFreshName Y.LocalNameKind + return + ( [ Y.Declare Y.TyBool y (Y.DeclareCopy (Y.BinOp Y.NotEqual (Y.callFunction "std::find" [] [Y.begin xs, Y.end xs, x]) (Y.end xs))) + ], + Y.Var y + ) + X.Sum -> go1' $ \xs -> do + y <- Y.newFreshName Y.LocalNameKind + return + ( [ Y.Declare Y.TyInt64 y (Y.DeclareCopy (Y.callFunction "std::accumulate" [] [Y.begin xs, Y.end xs, Y.litInt64 0])) + ], + Y.Var y + ) + X.ModSum -> go2' $ \xs m -> do + y <- Y.newFreshName Y.LocalNameKind + x <- Y.newFreshName Y.LocalNameKind + return + ( [ Y.Declare Y.TyInt64 y (Y.DeclareCopy (Y.litInt64 0)), + Y.ForEach + Y.TyInt64 + x + xs + [Y.Assign (Y.AssignExpr Y.AddAssign (Y.LeftVar y) (Y.callFunction "jikka::floormod" [] [Y.Var x, m]))] + ], + Y.callFunction "jikka::floormod" [] [Y.Var y, m] + ) + X.Product -> go1' $ \xs -> do + y <- Y.newFreshName Y.LocalNameKind + x <- Y.newFreshName Y.LocalNameKind + return + ( [ Y.Declare Y.TyInt64 y (Y.DeclareCopy (Y.litInt64 1)), + Y.ForEach + Y.TyInt64 + x + xs + [Y.Assign (Y.AssignExpr Y.MulAssign (Y.LeftVar y) (Y.Var x))] + ], + Y.Var y + ) + X.ModProduct -> go2' $ \xs m -> do + y <- Y.newFreshName Y.LocalNameKind + x <- Y.newFreshName Y.LocalNameKind + return + ( [ Y.Declare Y.TyInt64 y (Y.DeclareCopy (Y.litInt64 1)), + Y.ForEach + Y.TyInt64 + x + xs + [Y.Assign (Y.AssignExpr Y.SimpleAssign (Y.LeftVar y) (Y.callFunction "jikka::modmult" [] [Y.Var y, Y.Var x, m]))] + ], + Y.Var y + ) + X.Min1 t -> go1' $ \xs -> do + t <- runType t + y <- Y.newFreshName Y.LocalNameKind + return + ( [ Y.Declare t y (Y.DeclareCopy (Y.UnOp Y.Deref (Y.callFunction "std::min_element" [] [Y.begin xs, Y.end xs]))) + ], + Y.Var y + ) + X.Max1 t -> go1' $ \xs -> do + t <- runType t + y <- Y.newFreshName Y.LocalNameKind + return + ( [ Y.Declare t y (Y.DeclareCopy (Y.UnOp Y.Deref (Y.callFunction "std::max_element" [] [Y.begin xs, Y.end xs]))) + ], + Y.Var y + ) + X.ArgMin t -> go1' $ \xs -> do + t <- runType t + y <- Y.newFreshName Y.LocalNameKind + return + ( [ Y.Declare t y (Y.DeclareCopy (Y.BinOp Y.Sub (Y.callFunction "std::min_element" [] [Y.begin xs, Y.end xs]) (Y.begin xs))) + ], + Y.Var y + ) + X.ArgMax t -> go1' $ \xs -> do + t <- runType t + y <- Y.newFreshName Y.LocalNameKind + return + ( [ Y.Declare t y (Y.DeclareCopy (Y.BinOp Y.Sub (Y.callFunction "std::max_element" [] [Y.begin xs, Y.end xs]) (Y.begin xs))) + ], + Y.Var y + ) + X.All -> go1' $ \xs -> do + y <- Y.newFreshName Y.LocalNameKind + return + ( [ Y.Declare Y.TyBool y (Y.DeclareCopy (Y.BinOp Y.Equal (Y.callFunction "std::find" [] [Y.begin xs, Y.end xs, Y.Lit (Y.LitBool True)]) (Y.end xs))) + ], + Y.Var y + ) + X.Any -> go1' $ \xs -> do + y <- Y.newFreshName Y.LocalNameKind + return + ( [ Y.Declare Y.TyBool y (Y.DeclareCopy (Y.BinOp Y.NotEqual (Y.callFunction "std::find" [] [Y.begin xs, Y.end xs, Y.Lit (Y.LitBool False)]) (Y.end xs))) + ], + Y.Var y + ) + X.Sorted t -> go1' $ \xs -> do + t <- runType t + ys <- Y.newFreshName Y.LocalNameKind + return + ( [ Y.Declare (Y.TyVector t) ys (Y.DeclareCopy xs), + Y.callFunction' "std::sort" [] [Y.begin (Y.Var ys), Y.end (Y.Var ys)] + ], + Y.Var ys + ) + X.Reversed t -> go1' $ \xs -> do + t <- runType t + ys <- Y.newFreshName Y.LocalNameKind + return + ( [ Y.Declare (Y.TyVector t) ys (Y.DeclareCopy xs), + Y.callFunction' "std::reverse" [] [Y.begin (Y.Var ys), Y.end (Y.Var ys)] + ], + Y.Var ys + ) + X.Range1 -> go1 $ \n -> Y.Call Y.Range [n] + X.Range2 -> go2' $ \from to -> do + ys <- Y.newFreshName Y.LocalNameKind + return + ( [ Y.Declare (Y.TyVector Y.TyInt64) ys (Y.DeclareCopy (Y.vecCtor Y.TyInt64 [Y.BinOp Y.Sub to from])), + Y.callFunction' "std::iota" [] [Y.begin (Y.Var ys), Y.end (Y.Var ys), from] + ], + Y.Var ys + ) + X.Range3 -> go3' $ \from to step -> do + ys <- Y.newFreshName Y.LocalNameKind + i <- Y.newFreshName Y.LoopCounterNameKind + return + ( [ Y.Declare (Y.TyVector Y.TyInt64) ys Y.DeclareDefault, + Y.For + Y.TyInt32 + i + from + (Y.BinOp Y.LessThan (Y.Var i) to) + (Y.AssignExpr Y.AddAssign (Y.LeftVar i) step) + [ Y.callMethod' (Y.Var ys) "push_back" [Y.Var i] + ] + ], + Y.Var ys + ) + -- tuple functions + X.Tuple ts -> goN' $ \es -> do + ts <- mapM runType ts + return $ + if Y.shouldBeArray ts + then Y.Call (Y.ArrayExt (head ts)) es + else Y.Call (Y.StdTuple ts) es + X.Proj ts n -> go1' $ \e -> do + ts <- mapM runType ts + return . ([],) $ + if Y.shouldBeArray ts + then Y.at e (Y.Lit (Y.LitInt32 (fromIntegral n))) + else Y.Call (Y.StdGet (toInteger n)) [e] + -- comparison + X.LessThan _ -> go2 $ \e1 e2 -> Y.BinOp Y.LessThan e1 e2 + X.LessEqual _ -> go2 $ \e1 e2 -> Y.BinOp Y.LessEqual e1 e2 + X.GreaterThan _ -> go2 $ \e1 e2 -> Y.BinOp Y.GreaterThan e1 e2 + X.GreaterEqual _ -> go2 $ \e1 e2 -> Y.BinOp Y.GreaterEqual e1 e2 + X.Equal _ -> go2 $ \e1 e2 -> Y.BinOp Y.Equal e1 e2 + X.NotEqual _ -> go2 $ \e1 e2 -> Y.BinOp Y.NotEqual e1 e2 + -- combinational functions + X.Fact -> go1 $ \e -> Y.Call (Y.Function "jikka::fact" []) [e] + X.Choose -> go2 $ \e1 e2 -> Y.Call (Y.Function "jikka::choose" []) [e1, e2] + X.Permute -> go2 $ \e1 e2 -> Y.Call (Y.Function "jikka::permute" []) [e1, e2] + X.MultiChoose -> go2 $ \e1 e2 -> Y.Call (Y.Function "jikka::multichoose" []) [e1, e2] + -- data structures + X.ConvexHullTrickInit -> go0 $ Y.Call Y.ConvexHullTrickCtor [] + X.ConvexHullTrickGetMin -> go2 $ \cht x -> Y.Call (Y.Method "get_min") [cht, x] + X.ConvexHullTrickInsert -> go3 $ \cht a b -> Y.Call Y.ConvexHullTrickCopyAddLine [cht, a, b] + X.SegmentTreeInitList semigrp -> go1 $ \a -> Y.Call (Y.SegmentTreeCtor (runSemigroup semigrp)) [a] + X.SegmentTreeGetRange _ -> go3 $ \segtree l r -> Y.Call (Y.Method "prod") [segtree, l, r] + X.SegmentTreeSetPoint semigrp -> go3 $ \segtree i a -> Y.Call (Y.SegmentTreeCopySetPoint (runSemigroup semigrp)) [segtree, i, a] + +runExprFunction :: (MonadAlpha m, MonadError Error m) => Env -> X.Expr -> Y.Expr -> m ([Y.Statement], [Y.Statement], Y.Expr) +runExprFunction env f e = case f of + X.Lam x t body -> do + y <- renameVarName' Y.LocalArgumentNameKind x + (stmts, body) <- runExpr ((x, t, y) : env) body + let stmts' = map (Y.replaceStatement y e) stmts + let body' = Y.replaceExpr y e body + return ([], stmts', body') + f -> do + (stmts, f) <- runExpr env f + return (stmts, [], Y.CallExpr f [e]) + +runExprFunction2 :: (MonadAlpha m, MonadError Error m) => Env -> X.Expr -> Y.Expr -> Y.Expr -> m ([Y.Statement], [Y.Statement], Y.Expr) +runExprFunction2 env f e1 e2 = case f of + X.Lam2 x1 t1 x2 t2 body -> do + y1 <- renameVarName' Y.LocalArgumentNameKind x1 + y2 <- renameVarName' Y.LocalArgumentNameKind x2 + (stmts, body) <- runExpr ((x2, t2, y2) : (x1, t1, y1) : env) body + let stmts' = map (Y.replaceStatement y2 e2 . Y.replaceStatement y1 e1) stmts + let body' = Y.replaceExpr y2 e2 $ Y.replaceExpr y1 e1 body + return ([], stmts', body') + f -> do + (stmts, f) <- runExpr env f + return (stmts, [], Y.CallExpr (Y.CallExpr f [e1]) [e2]) + +runExpr :: (MonadAlpha m, MonadError Error m) => Env -> X.Expr -> m ([Y.Statement], Y.Expr) +runExpr env = \case + X.Var x -> do + y <- lookupVarName env x + return ([], Y.Var y) + X.Lit lit -> do + lit <- runLiteral env lit + return ([], lit) + e@(X.App _ _) -> do + let (f, args) = X.curryApp e + case f of + X.Lit (X.LitBuiltin builtin) -> do + let arity = arityOfBuiltin builtin + if length args < arity + then do + let (ts, ret) = X.uncurryFunTy (X.builtinToType builtin) + ts <- mapM runType ts + ret <- runType ret + xs <- replicateM (arity - length args) X.genVarName' + ys <- mapM (renameVarName' Y.LocalArgumentNameKind) xs + (stmts, e) <- runAppBuiltin env builtin (args ++ map X.Var xs) + let (_, e') = foldr (\(t, y) (ret, e) -> (Y.TyFunction ret [t], Y.Lam [(t, y)] ret [Y.Return e])) (ret, e) (zip (drop (length args) ts) ys) + return (stmts, e') + else + if length args == arity + then do + runAppBuiltin env builtin args + else do + (stmts, e) <- runAppBuiltin env builtin (take arity args) + args <- mapM (runExpr env) (drop arity args) + return (concatMap fst args ++ stmts, Y.CallExpr e (map snd args)) + _ -> do + args <- mapM (runExpr env) args + (stmts, f) <- runExpr env f + return (stmts ++ concatMap fst args, Y.CallExpr f (map snd args)) + e@(X.Lam _ _ _) -> do + let (args, body) = X.uncurryLam e + ys <- mapM (renameVarName' Y.LocalArgumentNameKind . fst) args + let env' = reverse (zipWith (\(x, t) y -> (x, t, y)) args ys) ++ env + ret <- runType =<< typecheckExpr env' body + (stmts, body) <- runExpr env' body + ts <- mapM (runType . snd) args + let (_, [Y.Return e]) = foldr (\(t, y) (ret, body) -> (Y.TyFunction ret [t], [Y.Return (Y.Lam [(t, y)] ret body)])) (ret, stmts ++ [Y.Return body]) (zip ts ys) + return ([], e) + X.Let x t e1 e2 -> do + y <- renameVarName' Y.LocalNameKind x + t' <- runType t + (stmts1, e1) <- runExpr env e1 + (stmts2, e2) <- runExpr ((x, t, y) : env) e2 + return (stmts1 ++ Y.Declare t' y (Y.DeclareCopy e1) : stmts2, e2) + +runToplevelFunDef :: (MonadAlpha m, MonadError Error m) => Env -> Y.VarName -> [(X.VarName, X.Type)] -> X.Type -> X.Expr -> m [Y.ToplevelStatement] +runToplevelFunDef env f args ret body = do + ret <- runType ret + args <- forM args $ \(x, t) -> do + y <- renameVarName' Y.ArgumentNameKind x + return (x, t, y) + (stmts, result) <- runExpr (reverse args ++ env) body + args <- forM args $ \(_, t, y) -> do + t <- runType t + return (t, y) + return [Y.FunDef ret f args (stmts ++ [Y.Return result])] + +runToplevelVarDef :: (MonadAlpha m, MonadError Error m) => Env -> Y.VarName -> X.Type -> X.Expr -> m [Y.ToplevelStatement] +runToplevelVarDef env x t e = do + t <- runType t + (stmts, e) <- runExpr env e + case stmts of + [] -> return [Y.VarDef t x e] + _ -> return [Y.VarDef t x (Y.CallExpr (Y.Lam [] t (stmts ++ [Y.Return e])) [])] + +runToplevelExpr :: (MonadAlpha m, MonadError Error m) => Env -> X.ToplevelExpr -> m [Y.ToplevelStatement] +runToplevelExpr env = \case + X.ResultExpr e -> do + t <- typecheckExpr env e + case X.uncurryFunTy t of + (ts@(_ : _), ret) -> do + let f = Y.VarName "solve" + (args, body) <- case X.uncurryLam e of + (args, body) | length args == length ts -> do + -- merge two sets of arguments which introduced by @FunTy@ and @Lam@ + args <- forM args $ \(x, t) -> do + y <- renameVarName' Y.ArgumentNameKind x + return (x, t, y) + (stmts, e) <- runExpr (reverse args ++ env) body + let body = stmts ++ [Y.Return e] + args' <- forM args $ \(_, t, y) -> do + t <- runType t + return (t, y) + return (args', body) + _ -> do + args <- forM ts $ \t -> do + t <- runType t + y <- Y.newFreshName Y.ArgumentNameKind + return (t, y) + (stmts, e) <- runExpr env e + let body = stmts ++ [Y.Return (Y.CallExpr e (map (Y.Var . snd) args))] + return (args, body) + ret <- runType ret + return [Y.FunDef ret f args body] + _ -> throwInternalError "solve function must be a function" -- TODO: add check in restricted Python + X.ToplevelLet x t e cont -> case (X.uncurryLam e, X.uncurryFunTy t) of + ((args@(_ : _), body), (ts@(_ : _), ret)) -> do + g <- renameVarName' Y.FunctionNameKind x + (args, body) <- + if length args < length ts + then do + xs <- replicateM (length ts - length args) X.genVarName' + let args' = args ++ zip xs (drop (length args) ts) + let body' = X.uncurryApp body (map X.Var xs) + return (args', body') + else return (args, body) + stmt <- runToplevelFunDef ((x, t, g) : env) g args ret body + cont <- runToplevelExpr ((x, t, g) : env) cont + return $ stmt ++ cont + _ -> do + y <- renameVarName' Y.ConstantNameKind x + stmt <- runToplevelVarDef env y t e + cont <- runToplevelExpr ((x, t, y) : env) cont + return $ stmt ++ cont + X.ToplevelLetRec f args ret body cont -> do + g <- renameVarName' Y.FunctionNameKind f + let t = X.curryFunTy (map snd args) ret + stmt <- runToplevelFunDef ((f, t, g) : env) g args ret body + cont <- runToplevelExpr ((f, t, g) : env) cont + return $ stmt ++ cont + +runProgram :: (MonadAlpha m, MonadError Error m) => X.Program -> m Y.Program +runProgram prog = Y.Program <$> runToplevelExpr [] prog + +run :: (MonadAlpha m, MonadError Error m) => X.Program -> m Y.Program +run prog = wrapError' "Jikka.CPlusPlus.Convert.FromCore" $ do + runProgram prog diff --git a/src/Jikka/CPlusPlus/Convert/MoveSemantics.hs b/src/Jikka/CPlusPlus/Convert/MoveSemantics.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/CPlusPlus/Convert/MoveSemantics.hs @@ -0,0 +1,185 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | +-- Module : Jikka.CPlusPlus.Convert.MoveSemantics +-- Description : removes unnecessary copying. / 無用なコピーを削除します。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.CPlusPlus.Convert.MoveSemantics + ( run, + ) +where + +import Control.Monad.State.Strict +import qualified Data.Map as M +import Data.Maybe +import qualified Data.Set as S +import Jikka.CPlusPlus.Language.Expr +import Jikka.CPlusPlus.Language.Util +import Jikka.CPlusPlus.Language.VariableAnalysis +import Jikka.Common.Error + +runExpr :: MonadState (M.Map VarName VarName) m => Expr -> m Expr +runExpr = \case + Var x -> do + y <- gets (M.lookup x) + return $ Var (fromMaybe x y) + Lit lit -> return $ Lit lit + UnOp op e -> UnOp op <$> runExpr e + BinOp op e1 e2 -> BinOp op <$> runExpr e1 <*> runExpr e2 + Cond e1 e2 e3 -> Cond <$> runExpr e1 <*> runExpr e2 <*> runExpr e3 + Lam args ret body -> Lam args ret <$> runStatements body [] + Call f args -> Call f <$> mapM runExpr args + CallExpr f args -> CallExpr <$> runExpr f <*> mapM runExpr args + +runLeftExpr :: MonadState (M.Map VarName VarName) m => LeftExpr -> m LeftExpr +runLeftExpr = \case + LeftVar x -> do + y <- gets (M.lookup x) + return $ LeftVar (fromMaybe x y) + LeftAt e1 e2 -> LeftAt <$> runLeftExpr e1 <*> runExpr e2 + LeftGet n e -> LeftGet n <$> runLeftExpr e + +runAssignExpr :: MonadState (M.Map VarName VarName) m => AssignExpr -> m AssignExpr +runAssignExpr = \case + AssignExpr op e1 e2 -> AssignExpr op <$> runLeftExpr e1 <*> runExpr e2 + AssignIncr e -> AssignIncr <$> runLeftExpr e + AssignDecr e -> AssignDecr <$> runLeftExpr e + +isMovable :: VarName -> [[Statement]] -> Bool +isMovable x cont = + let ReadWriteList rs _ = analyzeStatements (concat cont) + in x `S.notMember` rs + +runStatement :: MonadState (M.Map VarName VarName) m => Statement -> [[Statement]] -> m [Statement] +runStatement stmt cont = case stmt of + ExprStatement e -> do + e <- runExpr e + return [ExprStatement e] + Block stmts -> do + runStatements stmts cont + If e body1 body2 -> do + e <- runExpr e + body1 <- runStatements body1 cont + body2 <- traverse (`runStatements` cont) body2 + return [If e body1 body2] + For t x init pred incr body -> do + init <- runExpr init + pred <- runExpr pred + incr <- runAssignExpr incr + body <- runStatements body cont + return [For t x init pred incr body] + ForEach t x e body -> do + e <- runExpr e + body <- runStatements body cont + return [ForEach t x e body] + While e body -> do + e <- runExpr e + body <- runStatements body cont + return [While e body] + Declare t y init -> do + init <- case init of + DeclareDefault -> return DeclareDefault + DeclareCopy e -> DeclareCopy <$> runExpr e + DeclareInitialize es -> DeclareInitialize <$> mapM runExpr es + case init of + DeclareCopy (Var x) | x `isMovable` cont -> do + modify' (M.insert y x) + return [] + DeclareCopy (Call ConvexHullTrickCtor []) -> return [Declare t y DeclareDefault] + DeclareCopy (Call ConvexHullTrickCopyAddLine [Var x, a, b]) + | x `isMovable` cont -> do + modify' (M.insert y x) + return [callMethod' (Var x) "add_line" [a, b]] + DeclareCopy (Call (SegmentTreeCopySetPoint _) [Var x, i, a]) + | x `isMovable` cont -> do + modify' (M.insert y x) + return [callMethod' (Var x) "set" [i, a]] + _ -> do + return [Declare t y init] + DeclareDestructure xs e -> do + e <- runExpr e + return [DeclareDestructure xs e] + Assign e -> do + e <- runAssignExpr e + case e of + AssignExpr SimpleAssign (LeftVar y) (Var x) | x == y -> return [] + AssignExpr SimpleAssign (LeftVar y) (Call ConvexHullTrickCopyAddLine [Var x, a, b]) + | x == y -> return [callMethod' (Var x) "add_line" [a, b]] + | x `isMovable` cont -> do + modify' (M.insert y x) + return [callMethod' (Var x) "add_line" [a, b]] + | otherwise -> return [Assign e] + AssignExpr SimpleAssign (LeftVar y) (Call (SegmentTreeCopySetPoint _) [Var x, i, a]) + | x == y -> return [callMethod' (Var x) "set" [i, a]] + | x `isMovable` cont -> do + modify' (M.insert y x) + return [callMethod' (Var x) "set" [i, a]] + | otherwise -> return [Assign e] + _ -> return [Assign e] + Assert e -> do + e <- runExpr e + return [Assert e] + Return e -> do + e <- runExpr e + return [Return e] + +runStatements :: MonadState (M.Map VarName VarName) m => [Statement] -> [[Statement]] -> m [Statement] +runStatements stmts cont = case stmts of + [] -> return [] + stmt : stmts -> do + stmt <- runStatement stmt (stmts : cont) + stmts <- runStatements stmts cont + return (stmt ++ stmts) + +runToplevelStatement :: MonadState (M.Map VarName VarName) m => ToplevelStatement -> m ToplevelStatement +runToplevelStatement = \case + VarDef t x e -> VarDef t x <$> runExpr e + FunDef ret f args body -> FunDef ret f args <$> runStatements body [] + +runProgram :: Monad m => Program -> m Program +runProgram (Program decls) = (`evalStateT` M.empty) $ do + Program <$> mapM runToplevelStatement decls + +-- | `run` replaces superfluous copying. +-- +-- == Examples +-- +-- Before: +-- +-- > vector solve(vector a) { +-- > vector b = a; +-- > b[0] = 1; +-- > return b; +-- > } +-- +-- After: +-- +-- > vector solve(vector a) { +-- > a[0] = 1; +-- > return a; +-- > } +-- +-- Before: +-- +-- > int solve(int a, int b, int x) { +-- > jikka::convex_hull_trick cht = jikka::convex_hull_trick(); +-- > cht = jikka::convex_hull_trick::persistent_add_line(cht, a, b); +-- > return cht.get_min(x); +-- > } +-- +-- After: +-- +-- > int solve(int a, int b, int x) { +-- > jikka::convex_hull_trick cht; +-- > cht = cht.add_line(a, b); +-- > return cht.get_min(x); +-- > } +run :: MonadError Error m => Program -> m Program +run prog = wrapError' "Jikka.CPlusPlus.Convert.MoveSemantics" $ do + runProgram prog diff --git a/src/Jikka/CPlusPlus/Convert/OptimizeRange.hs b/src/Jikka/CPlusPlus/Convert/OptimizeRange.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/CPlusPlus/Convert/OptimizeRange.hs @@ -0,0 +1,60 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.CPlusPlus.Convert.OptimizeRange +-- Description : reduces about @range@ function. / @range@ 関数について簡約します。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.CPlusPlus.Convert.OptimizeRange + ( run, + ) +where + +import Jikka.CPlusPlus.Language.Expr +import Jikka.CPlusPlus.Language.Util +import Jikka.Common.Error + +runExpr :: Monad m => Expr -> m Expr +runExpr = \case + Call At [Call Range [_], i] -> return i + Call MethodSize [Call Range [n]] -> return n + e -> return e + +runStatement :: Monad m => Statement -> m Statement +runStatement = \case + ForEach _ x (Call Range [n]) body -> return $ repStatement x n body -- TODO: check n is not updated in body + stmt -> return stmt + +runProgram :: Monad m => Program -> m Program +runProgram = mapExprStatementProgramM runExpr runStatement + +-- | `run` replaces superfluous copying. +-- +-- == Examples +-- +-- Before: +-- +-- > int b = range(a).size(); +-- +-- After: +-- +-- > int b = a; +-- +-- Before: +-- +-- > for (int i : jikka::range(n)) { +-- > ... +-- > } +-- +-- After: +-- +-- > for (int i = 0; i < n; ++ i) { +-- > ... +-- > } +run :: MonadError Error m => Program -> m Program +run prog = wrapError' "Jikka.CPlusPlus.Convert.OptimizeRange" $ do + runProgram prog diff --git a/src/Jikka/CPlusPlus/Convert/UnpackTuples.hs b/src/Jikka/CPlusPlus/Convert/UnpackTuples.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/CPlusPlus/Convert/UnpackTuples.hs @@ -0,0 +1,217 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.CPlusPlus.Convert.UnpackTuples +-- Description : unpack tuples. / タプルを展開します。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.CPlusPlus.Convert.UnpackTuples + ( run, + ) +where + +import Control.Monad.State.Strict +import qualified Data.Map as M +import qualified Data.Set as S +import Jikka.CPlusPlus.Language.Expr +import Jikka.CPlusPlus.Language.Util +import Jikka.Common.Alpha +import Jikka.Common.Error + +runExpr :: (MonadAlpha m, MonadError Error m, MonadState (M.Map VarName [(Type, VarName)]) m) => Expr -> m Expr +runExpr = \case + Var x -> do + ys <- gets (M.lookup x) + return $ case ys of + Nothing -> Var x + Just ys -> + let es = map (Var . snd) ys + in if shouldBeArray (map fst ys) + then + let t = fst (head ys) + in Call (ArrayExt t) es + else + let ts = map fst ys + in Call (StdTuple ts) es + Lit lit -> return $ Lit lit + UnOp op e -> UnOp op <$> runExpr e + BinOp op e1 e2 -> BinOp op <$> runExpr e1 <*> runExpr e2 + Cond e1 e2 e3 -> Cond <$> runExpr e1 <*> runExpr e2 <*> runExpr e3 + Lam args ret body -> Lam args ret <$> runStatements body [] + Call f args -> runCall f args + CallExpr e args -> CallExpr <$> runExpr e <*> mapM runExpr args + +runCall :: (MonadAlpha m, MonadError Error m, MonadState (M.Map VarName [(Type, VarName)]) m) => Function -> [Expr] -> m Expr +runCall f args = do + args <- mapM runExpr args + case (f, args) of + (StdGet n, [Var x]) -> do + ys <- gets (M.lookup x) + case ys of + Just ys -> do + let es = map (Var . snd) ys + when (n < 0 || toInteger (length ys) <= n) $ do + throwInternalError "index out of range" + return $ es !! fromInteger n + Nothing -> return $ Call f args + (StdGet n, [Call (StdTuple _) es]) -> do + when (n < 0 || toInteger (length es) <= n) $ do + throwInternalError "index out of range" + return $ es !! fromInteger n + (At, [Var x, e2]) -> do + ys <- gets (M.lookup x) + case ys of + Just ys -> do + let es = map (Var . snd) ys + let n = case e2 of + Lit (LitInt32 n) -> Just n + Lit (LitInt64 n) -> Just n + _ -> Nothing + case n of + Just n -> do + when (n < 0 || toInteger (length ys) <= n) $ do + throwInternalError "index out of range" + return (es !! fromInteger n) + Nothing -> return $ Call f args + Nothing -> return $ Call f args + (At, [Call (ArrayExt _) es, e2]) -> do + let n = case e2 of + Lit (LitInt32 n) -> Just n + Lit (LitInt64 n) -> Just n + _ -> Nothing + case n of + Just n -> do + when (n < 0 || toInteger (length es) <= n) $ do + throwInternalError "index out of range" + return (es !! fromInteger n) + Nothing -> return $ Call f args + _ -> return $ Call f args + +runLeftExpr :: (MonadAlpha m, MonadError Error m, MonadState (M.Map VarName [(Type, VarName)]) m) => LeftExpr -> m LeftExpr +runLeftExpr = \case + LeftVar x -> return $ LeftVar x -- do nothing + LeftAt e1 e2 -> LeftAt <$> runLeftExpr e1 <*> runExpr e2 + LeftGet n e -> LeftGet n <$> runLeftExpr e + +runAssignExpr :: (MonadAlpha m, MonadError Error m, MonadState (M.Map VarName [(Type, VarName)]) m) => AssignExpr -> m AssignExpr +runAssignExpr = \case + AssignExpr op e1 e2 -> AssignExpr op <$> runLeftExpr e1 <*> runExpr e2 + AssignIncr e -> AssignIncr <$> runLeftExpr e + AssignDecr e -> AssignDecr <$> runLeftExpr e + +runStatement :: (MonadAlpha m, MonadError Error m, MonadState (M.Map VarName [(Type, VarName)]) m) => Statement -> [[Statement]] -> m [Statement] +runStatement stmt cont = case stmt of + ExprStatement e -> do + e <- runExpr e + return [ExprStatement e] + Block stmts -> do + runStatements stmts cont + If e body1 body2 -> do + e <- runExpr e + body1 <- runStatements body1 cont + body2 <- traverse (`runStatements` cont) body2 + return [If e body1 body2] + For t x init pred incr body -> do + init <- runExpr init + pred <- runExpr pred + incr <- runAssignExpr incr + body <- runStatements body cont + return [For t x init pred incr body] + ForEach t x e body -> do + e <- runExpr e + body <- runStatements body cont + return [ForEach t x e body] + While e body -> do + e <- runExpr e + body <- runStatements body cont + return [While e body] + Declare t x init -> do + init <- case init of + DeclareDefault -> return DeclareDefault + DeclareCopy e -> DeclareCopy <$> runExpr e + DeclareInitialize es -> DeclareInitialize <$> mapM runExpr es + case init of + DeclareCopy (Call (StdTuple ts) es) -> do + ys <- replicateM (length es) (renameVarName LocalNameKind (unVarName x)) + modify' (M.insert x (zip ts ys)) + return $ zipWith3 (\t y e -> Declare t y (DeclareCopy e)) ts ys es + DeclareCopy (Call (ArrayExt t) es) -> do + let ts = replicate (length es) t + ys <- replicateM (length es) (renameVarName LocalNameKind (unVarName x)) + modify' (M.insert x (zip ts ys)) + return $ zipWith3 (\t y e -> Declare t y (DeclareCopy e)) ts ys es + _ -> do + return [Declare t x init] + DeclareDestructure xs e -> do + e <- runExpr e + return [DeclareDestructure xs e] + Assign e -> do + e <- runAssignExpr e + case e of + AssignExpr SimpleAssign (LeftVar x) e -> do + ys <- gets (M.lookup x) + case ys of + Just ys -> do + let ts = map fst ys + let n = toInteger (length ts) + let es = case e of + Call (StdTuple _) es -> es + Call (ArrayExt _) es -> es + _ -> + if shouldBeArray ts + then map (\i -> Call At [e, litInt32 i]) [0 .. n - 1] + else map (\i -> Call (StdGet i) [e]) [0 .. n - 1] + return $ zipWith (\y e -> Assign (AssignExpr SimpleAssign (LeftVar y) e)) (map snd ys) es + Nothing -> return [Assign (AssignExpr SimpleAssign (LeftVar x) e)] + _ -> do + forM_ (S.toList (freeVarsAssignExpr e)) $ \x -> do + ys <- gets (M.lookup x) + case ys of + Just _ -> throwInternalError $ "wrong assignment to a tuple: " ++ unVarName x + Nothing -> return () + return [Assign e] + Assert e -> do + e <- runExpr e + return [Assert e] + Return e -> do + e <- runExpr e + return [Return e] + +runStatements :: (MonadAlpha m, MonadError Error m, MonadState (M.Map VarName [(Type, VarName)]) m) => [Statement] -> [[Statement]] -> m [Statement] +runStatements stmts cont = case stmts of + [] -> return [] + stmt : stmts -> do + stmt <- runStatement stmt (stmts : cont) + stmts <- runStatements stmts cont + return (stmt ++ stmts) + +runToplevelStatement :: (MonadAlpha m, MonadError Error m, MonadState (M.Map VarName [(Type, VarName)]) m) => ToplevelStatement -> m ToplevelStatement +runToplevelStatement = \case + VarDef t x e -> VarDef t x <$> runExpr e + FunDef ret f args body -> FunDef ret f args <$> runStatements body [] + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram (Program decls) = (`evalStateT` M.empty) $ do + Program <$> mapM runToplevelStatement decls + +-- | `run` unpack tuples. +-- +-- == Examples +-- +-- Before: +-- +-- > tuple c = make_tuple(a, b); +-- > func(get<0>(c), get<1>(c)); +-- +-- After: +-- +-- > int c0 = a; +-- > int c1 = b; +-- > func(c0, c1); +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.CPlusPlus.Convert.UnpackTuples" $ do + runProgram prog diff --git a/src/Jikka/CPlusPlus/Convert/UseInitialization.hs b/src/Jikka/CPlusPlus/Convert/UseInitialization.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/CPlusPlus/Convert/UseInitialization.hs @@ -0,0 +1,48 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.CPlusPlus.Convert.UseInitialization +-- Description : replaces declarations by assignments with initializations. / 代入による宣言を初期化による宣言で置き換えます。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.CPlusPlus.Convert.UseInitialization + ( run, + ) +where + +import Jikka.CPlusPlus.Language.Expr +import Jikka.CPlusPlus.Language.Util +import Jikka.Common.Alpha +import Jikka.Common.Error + +runStatement :: Statement -> Statement +runStatement = \case + Declare t x init -> case (t, init) of + (TyVector _, DeclareCopy (Call (VecCtor _) [])) -> Declare t x DeclareDefault + (TyVector _, DeclareCopy (Call (VecCtor _) es)) -> Declare t x (DeclareInitialize es) + (TyConvexHullTrick, DeclareCopy (Call ConvexHullTrickCtor es)) -> Declare t x (DeclareInitialize es) + (TySegmentTree _, DeclareCopy (Call (SegmentTreeCtor _) es)) -> Declare t x (DeclareInitialize es) + (_, _) -> Declare t x init + stmt -> stmt + +runProgram :: Program -> Program +runProgram = mapExprStatementProgram id runStatement + +-- | `run` unpack tuples. +-- +-- == Examples +-- +-- Before: +-- +-- > vector xs = vector(n, 0); +-- +-- After: +-- +-- > vector xs(n, 0); +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.CPlusPlus.Convert.UseInitialization" $ do + return $ runProgram prog diff --git a/src/Jikka/CPlusPlus/Format.hs b/src/Jikka/CPlusPlus/Format.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/CPlusPlus/Format.hs @@ -0,0 +1,351 @@ +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.CPlusPlus.Format +-- Description : converts the AST of C++ to strings. / C++ の抽象構文木を文字列に変換します。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- `Jikka.CPlusPlus.Format` module converts the AST for C++ to the plain source code. +module Jikka.CPlusPlus.Format + ( run, + run', + Code, + formatExpr, + formatType, + ) +where + +import Data.List (intercalate, isInfixOf) +import Data.Text (Text, pack) +import Jikka.CPlusPlus.Language.Expr +import Jikka.CPlusPlus.Language.Util +import Jikka.Common.Format.AutoIndent (makeIndentFromBraces) + +type Code = String + +-- | +data Prec + = IdentPrec + | ScopeResolutionPrec + | -- | the precidense of function calls + FunCallPrec + | -- | the precidense of @!@ and @~@ + UnaryPrec + | PointerToMemberPrec + | -- | the precidense of @*@, @/@, and @%@ + MultPrec + | -- | the precidense of @+@, @-@ + AddPrec + | -- | the precidense of @<<@, @>>@ + ShiftPrec + | -- | the precidense of @<@, @<=@, @>@, and @>=@ + LessThanPrec + | -- | the precidense of @==@ and @!=@ + EqualPrec + | -- | the precidense of @&@ + BitAndPrec + | -- | the precidense of @^@ + BitXorPrec + | -- | the precidense of @|@ + BitOrPrec + | -- | the precidense of @&&@ + AndPrec + | -- | the precidense of @||@ + OrPrec + | -- | the precidense of the conditional operator @?@ and @:@ + CondPrec + | -- | the precidense of the conditional operator @=@, @+=@, @-=@, ... + AssignPrec + | ThrowPrec + | -- | the precidense of @,@ + CommaPrec + | ParenPrec + deriving (Eq, Ord, Show, Read) + +data Assoc + = NoAssoc + | LeftToRight + | RightToLeft + deriving (Eq, Ord, Show, Read) + +assocOf :: Prec -> Assoc +assocOf = \case + IdentPrec -> NoAssoc + ScopeResolutionPrec -> NoAssoc + FunCallPrec -> LeftToRight + UnaryPrec -> RightToLeft + PointerToMemberPrec -> LeftToRight + MultPrec -> LeftToRight + AddPrec -> LeftToRight + ShiftPrec -> LeftToRight + LessThanPrec -> LeftToRight + EqualPrec -> LeftToRight + BitAndPrec -> LeftToRight + BitXorPrec -> LeftToRight + BitOrPrec -> LeftToRight + AndPrec -> LeftToRight + OrPrec -> LeftToRight + CondPrec -> RightToLeft + AssignPrec -> RightToLeft + ThrowPrec -> RightToLeft + CommaPrec -> LeftToRight + ParenPrec -> NoAssoc + +formatUnaryOp :: UnaryOp -> (Code, Prec) +formatUnaryOp = \case + IntNop -> ("+", UnaryPrec) + Negate -> ("-", UnaryPrec) + BitNot -> ("~", UnaryPrec) + Not -> ("not", UnaryPrec) + Deref -> ("*", UnaryPrec) + +formatBinaryOp :: BinaryOp -> (Code, Prec) +formatBinaryOp = \case + Add -> ("+", AddPrec) + Sub -> ("-", AddPrec) + Mul -> ("*", MultPrec) + Div -> ("/", MultPrec) + Mod -> ("%", MultPrec) + BitLeftShift -> ("<<", ShiftPrec) + BitRightShift -> (">>", ShiftPrec) + LessThan -> ("<", LessThanPrec) + LessEqual -> ("<=", LessThanPrec) + GreaterThan -> (">", LessThanPrec) + GreaterEqual -> (">=", LessThanPrec) + Equal -> ("==", EqualPrec) + NotEqual -> ("!=", EqualPrec) + BitAnd -> ("&", BitAndPrec) + BitXor -> ("^", BitXorPrec) + BitOr -> ("|", BitOrPrec) + And -> ("and", AndPrec) + Or -> ("or", OrPrec) + +formatAssignOp :: AssignOp -> (Code, Prec) +formatAssignOp = \case + SimpleAssign -> ("=", AssignPrec) + AddAssign -> ("+=", AssignPrec) + SubAssign -> ("-=", AssignPrec) + MulAssign -> ("*=", AssignPrec) + DivAssign -> ("/=", AssignPrec) + ModAssign -> ("%=", AssignPrec) + BitLeftShiftAssign -> ("<<=", AssignPrec) + BitRightShiftAssign -> (">>=", AssignPrec) + BitAndAssign -> ("&=", AssignPrec) + BitOrAssign -> ("|=", AssignPrec) + BitXorAssign -> ("^=", AssignPrec) + +-- | `resolvePrec` inserts parens to the given string if required. +-- +-- >>> resolvePrec MultPrec ("1 + 2", AddPrec) ++ " * 3" +-- "(1 + 2) * 3" +-- +-- >>> resolvePrec AddPrec ("1 * 2", MultPrec) ++ " + 3" +-- "1 * 2 + 3" +-- +-- >>> resolvePrec CommaPrec ("1, 2", CommaPrec) ++ ", 3" +-- "1, 2, 3" +resolvePrec :: Prec -> (Code, Prec) -> Code +resolvePrec cur (s, prv) + | cur < prv = "(" ++ s ++ ")" + | otherwise = s + +-- | `resolvePrecLeft` inserts parens to the given string if required. +-- +-- >>> resolvePrecLeft AddPrec ("1 - 2", AddPrec) ++ " - 3" +-- "1 - 2 - 3" +resolvePrecLeft :: Prec -> (Code, Prec) -> Code +resolvePrecLeft cur (s, prv) + | cur < prv || (cur == prv && assocOf cur /= LeftToRight) = "(" ++ s ++ ")" + | otherwise = s + +-- | `resolvePrecRight` inserts parens to the given string if required. +-- +-- >>> "1 - " ++ resolvePrecRight AddPrec ("2 - 3", AddPrec) +-- "1 - (2 - 3)" +resolvePrecRight :: Prec -> (Code, Prec) -> Code +resolvePrecRight cur (s, prv) + | cur < prv || (cur == prv && assocOf cur /= RightToLeft) = "(" ++ s ++ ")" + | otherwise = s + +formatType :: Type -> Code +formatType = \case + TyAuto -> "auto" + TyVoid -> "void" + TyInt -> "int" + TyInt32 -> "int32_t" + TyInt64 -> "int64_t" + TyBool -> "bool" + TyTuple ts -> "std::tuple<" ++ intercalate ", " (map formatType ts) ++ ">" + TyVector t -> "std::vector<" ++ formatType t ++ ">" + TyArray t n -> "std::array<" ++ formatType t ++ ", " ++ show n ++ ">" + TyString -> "std::string" + TyFunction t ts -> "std::function<" ++ formatType t ++ " (" ++ intercalate ", " (map formatType ts) ++ ")>" + TyConvexHullTrick -> "jikka::convex_hull_trick" + TySegmentTree mon -> case mon of + MonoidIntPlus -> "atcoder::segtree" + MonoidIntMin -> "atcoder::segtree" + MonoidIntMax -> "atcoder::segtree" + TyIntValue n -> show n + +formatLiteral :: Literal -> Code +formatLiteral = \case + LitInt32 n -> show n + LitInt64 n -> show n + LitBool p -> if p then "true" else "false" + LitChar c -> show c + LitString s -> show s + +formatExpr' :: Prec -> Expr -> Code +formatExpr' prec = resolvePrec prec . formatExpr + +formatExpr :: Expr -> (Code, Prec) +formatExpr = \case + Var x -> (unVarName x, IdentPrec) + Lit lit -> (formatLiteral lit, IdentPrec) + UnOp op e -> + let (op', prec) = formatUnaryOp op + e' = formatExpr' prec e + in (op' ++ " " ++ e', prec) + BinOp op e1 e2 -> + let (op', prec) = formatBinaryOp op + e1' = resolvePrecLeft prec (formatExpr e1) + e2' = resolvePrecRight prec (formatExpr e2) + in (e1' ++ " " ++ op' ++ " " ++ e2', prec) + Lam args ret body -> + let args' = map (\(t, x) -> formatType t ++ " " ++ unVarName x) args + ret' = formatType ret + body' = concatMap formatStatement body + in ("[=](" ++ intercalate ", " args' ++ ") -> " ++ ret' ++ "{ " ++ unwords body' ++ " }", FunCallPrec) + Call f args -> + let args' = intercalate ", " (map (formatExpr' CommaPrec) args) + call f = (f ++ "(" ++ args' ++ ")", FunCallPrec) + method f = case args of + [] -> error $ "Jikka.CPlusPlus.Language.Format.formatExpr: no receiver for method: " ++ f + e : args -> (formatExpr' FunCallPrec e ++ "." ++ f ++ "(" ++ intercalate ", " (map (formatExpr' CommaPrec) args) ++ ")", FunCallPrec) + in case f of + Function f ts -> call $ unFunName f ++ (if null ts then "" else "<" ++ intercalate ", " (map formatType ts) ++ ">") + Method f -> method $ unFunName f + At -> case args of + [e1, e2] -> + let e1' = formatExpr' FunCallPrec e1 + e2' = formatExpr' FunCallPrec e2 + in (e1' ++ "[" ++ e2' ++ "]", FunCallPrec) + _ -> error $ "Jikka.CPlusPlus.Language.Format.formatExpr: wrong number of arguments for subscription: " ++ show (length args) + Cast t -> call $ formatType t + StdTuple ts -> call $ "std::tuple<" ++ intercalate ", " (map formatType ts) ++ ">" + StdGet n -> call $ "std::get<" ++ show n ++ ">" + ArrayExt t -> ("std::array<" ++ formatType t ++ ", " ++ show (length args) ++ ">{" ++ args' ++ "}", IdentPrec) + VecExt t -> ("std::vector<" ++ formatType t ++ ">{" ++ args' ++ "}", IdentPrec) + VecCtor t -> call $ "std::vector<" ++ formatType t ++ ">" + Range -> call "jikka::range" + MethodSize -> method "size" + ConvexHullTrickCtor -> call "jikka::convex_hull_trick" + ConvexHullTrickCopyAddLine -> call "jikka::convex_hull_trick::add_line" + SegmentTreeCtor mon -> call (formatType (TySegmentTree mon)) + SegmentTreeCopySetPoint _ -> call "jikka::segment_tree_set" + CallExpr f args -> + let f' = formatExpr' FunCallPrec f + args' = intercalate ", " (map (formatExpr' CommaPrec) args) + in (f' ++ "(" ++ args' ++ ")", FunCallPrec) + Cond e1 e2 e3 -> + let e1' = resolvePrecLeft CondPrec (formatExpr e1) + e2' = resolvePrec CondPrec (formatExpr e2) + e3' = resolvePrecRight CondPrec (formatExpr e3) + in (e1' ++ " ? " ++ e2' ++ " : " ++ e3', CondPrec) + +formatLeftExpr :: LeftExpr -> (Code, Prec) +formatLeftExpr = formatExpr . fromLeftExpr + +formatAssignExpr :: AssignExpr -> (Code, Prec) +formatAssignExpr = \case + AssignExpr op e1 e2 -> + let (op', prec) = formatAssignOp op + e1' = resolvePrecLeft prec (formatLeftExpr e1) + e2' = resolvePrecRight prec (formatExpr e2) + in (e1' ++ " " ++ op' ++ " " ++ e2', AssignPrec) + AssignIncr e -> ("++ " ++ resolvePrec UnaryPrec (formatLeftExpr e), UnaryPrec) + AssignDecr e -> ("-- " ++ resolvePrec UnaryPrec (formatLeftExpr e), UnaryPrec) + +formatStatement :: Statement -> [Code] +formatStatement = \case + ExprStatement e -> [formatExpr' ParenPrec e ++ ";"] + Block stmts -> ["{"] ++ concatMap formatStatement stmts ++ ["}"] + If e body1 body2 -> + let e' = formatExpr' ParenPrec e + body1' = concatMap formatStatement body1 + in case body2 of + Nothing -> ["if (" ++ e' ++ ") {"] ++ body1' ++ ["}"] + Just body2 -> case concatMap formatStatement body2 of + (('i' : 'f' : ' ' : '(' : line) : lines) -> ["if (" ++ e' ++ ") {"] ++ body1' ++ ["} else if (" ++ line] ++ lines + body2 -> ["if (" ++ e' ++ ") {"] ++ body1' ++ ["} else {"] ++ body2 ++ ["}"] + For t x init cond incr body -> + let t' = formatType t + init' = formatExpr' ParenPrec init + cond' = formatExpr' ParenPrec cond + incr' = resolvePrec ParenPrec $ formatAssignExpr incr + body' = concatMap formatStatement body + in ["for (" ++ t' ++ " " ++ unVarName x ++ " = " ++ init' ++ "; " ++ cond' ++ "; " ++ incr' ++ ") {"] ++ body' ++ ["}"] + ForEach t x xs body -> + let t' = formatType t + xs' = formatExpr' ParenPrec xs + body' = concatMap formatStatement body + in ["for (" ++ t' ++ " " ++ unVarName x ++ " : " ++ xs' ++ ") {"] ++ body' ++ ["}"] + While cond body -> + let cond' = formatExpr' ParenPrec cond + body' = concatMap formatStatement body + in ["while (" ++ cond' ++ ") {"] ++ body' ++ ["}"] + Declare t x init -> + let t' = formatType t + init' = case init of + DeclareDefault -> "" + DeclareCopy e -> " = " ++ resolvePrecRight AssignPrec (formatExpr e) + DeclareInitialize es -> "(" ++ intercalate ", " (map (formatExpr' CommaPrec) es) ++ ")" + in [t' ++ " " ++ unVarName x ++ init' ++ ";"] + DeclareDestructure xs e -> ["auto [" ++ intercalate ", " (map unVarName xs) ++ "] = " ++ resolvePrecRight AssignPrec (formatExpr e) ++ ";"] + Assign e -> [resolvePrec ParenPrec (formatAssignExpr e) ++ ";"] + Assert e -> ["assert (" ++ formatExpr' ParenPrec e ++ ");"] + Return e -> ["return " ++ formatExpr' ParenPrec e ++ ";"] + +formatToplevelStatement :: ToplevelStatement -> [Code] +formatToplevelStatement = \case + VarDef t x e -> [formatType t ++ " " ++ unVarName x ++ " = " ++ resolvePrecRight AssignPrec (formatExpr e) ++ ";"] + FunDef ret f args body -> + let ret' = formatType ret + args' = intercalate ", " $ map (\(t, x) -> formatType t ++ " " ++ unVarName x) args + body' = concatMap formatStatement body + in [ret' ++ " " ++ unVarName f ++ "(" ++ args' ++ ") {"] ++ body' ++ ["}"] + +formatProgram :: Program -> [Code] +formatProgram prog = + let body = concatMap formatToplevelStatement (decls prog) + standardHeaders = + [ "#include ", + "#include ", + "#include ", + "#include ", + "#include ", + "#include ", + "#include ", + "#include ", + "#include " + ] + additionalHeader = + map snd $ + filter + (\(key, _) -> key `isInfixOf` unlines body) + [ ("jikka::", "#include \"jikka/base.hpp\""), + ("jikka::convex_hull_trick", "#include \"jikka/convex_hull_trick.hpp\""), + ("atcoder::segtree", "#include \"jikka/segment_tree.hpp\""), + ("atcoder::segtree", "#include ") + ] + in standardHeaders ++ additionalHeader ++ body + +run' :: Program -> String +run' = unlines . makeIndentFromBraces 4 . formatProgram + +run :: Applicative m => Program -> m Text +run = pure . pack . run' diff --git a/src/Jikka/CPlusPlus/Language/Expr.hs b/src/Jikka/CPlusPlus/Language/Expr.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/CPlusPlus/Language/Expr.hs @@ -0,0 +1,168 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} + +-- | +-- Module : Jikka.CPlusPlus.Language.Expr +-- Description : contains data types of C++ language. / C++ のためのデータ型を含みます。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- `Jikka.CPlusPlus.Language.Expr` module has the basic data types for C++ language. +-- The data types are intended to use for the code generation. +module Jikka.CPlusPlus.Language.Expr where + +import Data.String (IsString) + +newtype VarName = VarName {unVarName :: String} deriving (Eq, Ord, Show, Read, IsString) + +newtype FunName = FunName {unFunName :: String} deriving (Eq, Ord, Show, Read, IsString) + +data Type + = TyAuto + | TyVoid + | TyBool + | TyInt + | TyInt32 + | TyInt64 + | TyTuple [Type] + | TyVector Type + | TyArray Type Integer + | TyString + | TyFunction Type [Type] + | TyConvexHullTrick + | TySegmentTree Monoid' + | -- | for template parameters + TyIntValue Integer + deriving (Eq, Ord, Show, Read) + +data Monoid' + = MonoidIntPlus + | MonoidIntMin + | MonoidIntMax + deriving (Eq, Ord, Show, Read) + +data Literal + = LitInt32 Integer + | LitInt64 Integer + | LitBool Bool + | LitChar Char + | LitString String + deriving (Eq, Ord, Show, Read) + +data Function + = Function FunName [Type] + | Method FunName + | At + | Cast Type + | StdTuple [Type] + | StdGet Integer + | ArrayExt Type + | VecExt Type + | VecCtor Type + | Range + | MethodSize + | ConvexHullTrickCtor + | ConvexHullTrickCopyAddLine + | SegmentTreeCtor Monoid' + | SegmentTreeCopySetPoint Monoid' + deriving (Eq, Ord, Show, Read) + +data UnaryOp + = IntNop + | Negate + | BitNot + | Not + | Deref + deriving (Eq, Ord, Show, Read) + +data BinaryOp + = Add + | Sub + | Mul + | Div + | Mod + | BitAnd + | BitOr + | BitXor + | BitLeftShift + | BitRightShift + | And + | Or + | LessThan + | LessEqual + | GreaterThan + | GreaterEqual + | Equal + | NotEqual + deriving (Eq, Ord, Show, Read) + +data AssignOp + = SimpleAssign + | AddAssign + | SubAssign + | MulAssign + | DivAssign + | ModAssign + | BitLeftShiftAssign + | BitRightShiftAssign + | BitAndAssign + | BitOrAssign + | BitXorAssign + deriving (Eq, Ord, Show, Read) + +data Expr + = Var VarName + | Lit Literal + | UnOp UnaryOp Expr + | BinOp BinaryOp Expr Expr + | Cond Expr Expr Expr + | Lam [(Type, VarName)] Type [Statement] + | Call Function [Expr] + | CallExpr Expr [Expr] + deriving (Eq, Ord, Show, Read) + +data LeftExpr + = LeftVar VarName + | LeftAt LeftExpr Expr + | -- | @std::get@ + LeftGet Integer LeftExpr + deriving (Eq, Ord, Show, Read) + +data AssignExpr + = AssignExpr AssignOp LeftExpr Expr + | AssignIncr LeftExpr + | AssignDecr LeftExpr + deriving (Eq, Ord, Show, Read) + +data DeclareRight + = DeclareDefault + | DeclareCopy Expr + | -- | This is only for better formatting. This should not be used while optimization phases. + DeclareInitialize [Expr] + deriving (Eq, Ord, Show, Read) + +data Statement + = ExprStatement Expr + | Block [Statement] + | If Expr [Statement] (Maybe [Statement]) + | For Type VarName Expr Expr AssignExpr [Statement] + | ForEach Type VarName Expr [Statement] + | While Expr [Statement] + | Declare Type VarName DeclareRight + | DeclareDestructure [VarName] Expr + | Assign AssignExpr + | Assert Expr + | Return Expr + deriving (Eq, Ord, Show, Read) + +data ToplevelStatement + = VarDef Type VarName Expr + | FunDef Type VarName [(Type, VarName)] [Statement] + deriving (Eq, Ord, Show, Read) + +newtype Program = Program + { decls :: [ToplevelStatement] + } + deriving (Eq, Ord, Show, Read) diff --git a/src/Jikka/CPlusPlus/Language/Util.hs b/src/Jikka/CPlusPlus/Language/Util.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/CPlusPlus/Language/Util.hs @@ -0,0 +1,220 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.CPlusPlus.Language.Util where + +import Control.Monad.Identity +import Data.Char (isAlphaNum) +import qualified Data.Set as S +import Jikka.CPlusPlus.Language.Expr +import Jikka.Common.Alpha + +fromLeftExpr :: LeftExpr -> Expr +fromLeftExpr = \case + LeftVar x -> Var x + LeftAt x e -> Call At [fromLeftExpr x, e] + LeftGet n e -> Call (Function "std::get" [TyIntValue n]) [fromLeftExpr e] + +data NameKind + = LocalNameKind + | LocalArgumentNameKind + | LoopCounterNameKind + | ConstantNameKind + | FunctionNameKind + | ArgumentNameKind + deriving (Eq, Ord, Show, Read) + +fromNameKind :: NameKind -> String +fromNameKind = \case + LocalNameKind -> "x" + LocalArgumentNameKind -> "b" + LoopCounterNameKind -> "i" + ConstantNameKind -> "c" + FunctionNameKind -> "f" + ArgumentNameKind -> "a" + +newFreshName :: MonadAlpha m => NameKind -> m VarName +newFreshName kind = renameVarName kind "" + +renameVarName :: MonadAlpha m => NameKind -> String -> m VarName +renameVarName kind hint = do + i <- nextCounter + let prefix = case takeWhile (\c -> isAlphaNum c || c == '_') hint of + "" -> fromNameKind kind + hint' -> hint' ++ "_" + return (VarName (prefix ++ show i)) + +freeVars :: Expr -> S.Set VarName +freeVars = \case + Var x -> S.singleton x + Lit _ -> S.empty + UnOp _ e -> freeVars e + BinOp _ e1 e2 -> freeVars e1 <> freeVars e2 + Cond e1 e2 e3 -> freeVars e1 <> freeVars e2 <> freeVars e3 + Lam args _ body -> freeVarsStatements body S.\\ S.fromList (map snd args) + Call _ args -> mconcat (map freeVars args) + CallExpr f args -> freeVars f <> mconcat (map freeVars args) + +freeVarsStatements :: [Statement] -> S.Set VarName +freeVarsStatements = mconcat . map freeVarsStatement + +freeVarsStatement :: Statement -> S.Set VarName +freeVarsStatement = \case + ExprStatement e -> freeVars e + Block stmts -> freeVarsStatements stmts + If e body1 body2 -> freeVars e <> freeVarsStatements body1 <> S.unions (fmap freeVarsStatements body2) + For _ x init pred incr body -> S.singleton x <> freeVars init <> freeVars pred <> freeVarsAssignExpr incr <> freeVarsStatements body + ForEach _ x e body -> S.singleton x <> freeVars e <> freeVarsStatements body + While e body -> freeVars e <> freeVarsStatements body + Declare _ x init -> S.singleton x <> freeVarsDeclareRight init + DeclareDestructure xs e -> S.fromList xs <> freeVars e + Assign e -> freeVarsAssignExpr e + Assert e -> freeVars e + Return e -> freeVars e + +freeVarsDeclareRight :: DeclareRight -> S.Set VarName +freeVarsDeclareRight = \case + DeclareDefault -> S.empty + DeclareCopy e -> freeVars e + DeclareInitialize es -> S.unions (map freeVars es) + +freeVarsAssignExpr :: AssignExpr -> S.Set VarName +freeVarsAssignExpr = \case + AssignExpr _ e1 e2 -> freeVarsLeftExpr e1 <> freeVars e2 + AssignIncr e -> freeVarsLeftExpr e + AssignDecr e -> freeVarsLeftExpr e + +freeVarsLeftExpr :: LeftExpr -> S.Set VarName +freeVarsLeftExpr = \case + LeftVar x -> S.singleton x + LeftAt e1 e2 -> freeVarsLeftExpr e1 <> freeVars e2 + LeftGet _ e -> freeVarsLeftExpr e + +shouldBeArray :: [Type] -> Bool +shouldBeArray ts = not (null ts) && ts == replicate (length ts) (head ts) + +cinStatement :: Expr -> Statement +cinStatement e = ExprStatement (BinOp BitRightShift (Var "std::cin") e) + +coutStatement :: Expr -> Statement +coutStatement e = ExprStatement (BinOp BitLeftShift (BinOp BitLeftShift (Var "std::cout") e) (Lit (LitChar ' '))) + +repStatement :: VarName -> Expr -> [Statement] -> Statement +repStatement i n body = For TyInt32 i (Lit (LitInt32 0)) (BinOp LessThan (Var i) n) (AssignIncr (LeftVar i)) body + +litInt64 :: Integer -> Expr +litInt64 n = Lit (LitInt64 n) + +litInt32 :: Integer -> Expr +litInt32 n = Lit (LitInt32 n) + +incrExpr :: Expr -> Expr +incrExpr e = BinOp Add e (Lit (LitInt32 1)) + +size :: Expr -> Expr +size e = Call MethodSize [e] + +at :: Expr -> Expr -> Expr +at e i = Call At [e, i] + +cast :: Type -> Expr -> Expr +cast t e = Call (Cast t) [e] + +assignSimple :: VarName -> Expr -> Statement +assignSimple x e = Assign (AssignExpr SimpleAssign (LeftVar x) e) + +assignAt :: VarName -> Expr -> Expr -> Statement +assignAt xs i e = Assign (AssignExpr SimpleAssign (LeftAt (LeftVar xs) i) e) + +callFunction :: FunName -> [Type] -> [Expr] -> Expr +callFunction f ts args = Call (Function f ts) args + +callFunction' :: FunName -> [Type] -> [Expr] -> Statement +callFunction' = ((ExprStatement .) .) . callFunction + +callMethod :: Expr -> FunName -> [Expr] -> Expr +callMethod e f args = Call (Method f) (e : args) + +callMethod' :: Expr -> FunName -> [Expr] -> Statement +callMethod' = ((ExprStatement .) .) . callMethod + +vecCtor :: Type -> [Expr] -> Expr +vecCtor t es = Call (VecCtor t) es + +begin :: Expr -> Expr +begin e = Call (Method "begin") [e] + +end :: Expr -> Expr +end e = Call (Method "end") [e] + +mapExprStatementExprM :: Monad m => (Expr -> m Expr) -> (Statement -> m Statement) -> Expr -> m Expr +mapExprStatementExprM f g = go + where + go = \case + Var x -> f (Var x) + Lit lit -> f (Lit lit) + UnOp op e -> f . UnOp op =<< go e + BinOp op e1 e2 -> f =<< (BinOp op <$> go e1 <*> go e2) + Cond e1 e2 e3 -> f =<< (Cond <$> go e1 <*> go e2 <*> go e3) + Lam args ret body -> f . Lam args ret =<< mapM (mapExprStatementStatementM f g) body + Call g args -> f . Call g =<< mapM go args + CallExpr g args -> f =<< (CallExpr <$> go g <*> mapM go args) + +mapExprStatementLeftExprM :: Monad m => (Expr -> m Expr) -> (Statement -> m Statement) -> LeftExpr -> m LeftExpr +mapExprStatementLeftExprM f g = \case + LeftVar x -> return $ LeftVar x + LeftAt e1 e2 -> LeftAt <$> mapExprStatementLeftExprM f g e1 <*> mapExprStatementExprM f g e2 + LeftGet n e -> LeftGet n <$> mapExprStatementLeftExprM f g e + +mapExprStatementAssignExprM :: Monad m => (Expr -> m Expr) -> (Statement -> m Statement) -> AssignExpr -> m AssignExpr +mapExprStatementAssignExprM f g = \case + AssignExpr op e1 e2 -> AssignExpr op <$> mapExprStatementLeftExprM f g e1 <*> mapExprStatementExprM f g e2 + AssignIncr e -> AssignIncr <$> mapExprStatementLeftExprM f g e + AssignDecr e -> AssignDecr <$> mapExprStatementLeftExprM f g e + +mapExprStatementStatementM :: Monad m => (Expr -> m Expr) -> (Statement -> m Statement) -> Statement -> m Statement +mapExprStatementStatementM f g = go + where + go' = mapExprStatementExprM f g + go = \case + ExprStatement e -> g . ExprStatement =<< go' e + Block stmts -> g . Block =<< mapM go stmts + If e body1 body2 -> g =<< (If <$> go' e <*> mapM go body1 <*> traverse (mapM go) body2) + For t x init pred incr body -> g =<< (For t x <$> go' init <*> go' pred <*> mapExprStatementAssignExprM f g incr <*> mapM go body) + ForEach t x e body -> g =<< (ForEach t x <$> go' e <*> mapM go body) + While e body -> g =<< (While <$> go' e <*> mapM go body) + Declare t x init -> do + init <- case init of + DeclareDefault -> return DeclareDefault + DeclareCopy e -> DeclareCopy <$> go' e + DeclareInitialize es -> DeclareInitialize <$> mapM go' es + g $ Declare t x init + DeclareDestructure xs e -> g . DeclareDestructure xs =<< go' e + Assign e -> g . Assign =<< mapExprStatementAssignExprM f g e + Assert e -> g . Assert =<< go' e + Return e -> g . Return =<< go' e + +mapExprStatementToplevelStatementM :: Monad m => (Expr -> m Expr) -> (Statement -> m Statement) -> ToplevelStatement -> m ToplevelStatement +mapExprStatementToplevelStatementM f g = \case + VarDef t x e -> VarDef t x <$> mapExprStatementExprM f g e + FunDef ret h args body -> FunDef ret h args <$> mapM (mapExprStatementStatementM f g) body + +mapExprStatementProgramM :: Monad m => (Expr -> m Expr) -> (Statement -> m Statement) -> Program -> m Program +mapExprStatementProgramM f g (Program decls) = Program <$> mapM (mapExprStatementToplevelStatementM f g) decls + +mapExprStatementProgram :: (Expr -> Expr) -> (Statement -> Statement) -> Program -> Program +mapExprStatementProgram f g = runIdentity . mapExprStatementProgramM (return . f) (return . g) + +replaceExpr :: VarName -> Expr -> Expr -> Expr +replaceExpr x e = runIdentity . mapExprStatementExprM go return + where + go = \case + Var y | y == x -> return e + e' -> return e' + +replaceStatement :: VarName -> Expr -> Statement -> Statement +replaceStatement x e = runIdentity . mapExprStatementStatementM go return + where + go = \case + Var y | y == x -> return e + e' -> return e' diff --git a/src/Jikka/CPlusPlus/Language/VariableAnalysis.hs b/src/Jikka/CPlusPlus/Language/VariableAnalysis.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/CPlusPlus/Language/VariableAnalysis.hs @@ -0,0 +1,71 @@ +{-# LANGUAGE LambdaCase #-} + +module Jikka.CPlusPlus.Language.VariableAnalysis where + +import qualified Data.Set as S +import Jikka.CPlusPlus.Language.Expr + +data ReadWriteList = ReadWriteList + { readList :: S.Set VarName, + writeList :: S.Set VarName + } + deriving (Eq, Ord, Show, Read) + +instance Semigroup ReadWriteList where + ReadWriteList rs ws <> ReadWriteList rs' ws' = ReadWriteList (rs <> rs') (ws <> ws') + +instance Monoid ReadWriteList where + mempty = ReadWriteList S.empty S.empty + +readVariable :: VarName -> ReadWriteList +readVariable x = ReadWriteList (S.singleton x) S.empty + +writeVariable :: VarName -> ReadWriteList +writeVariable x = ReadWriteList S.empty (S.singleton x) + +analyzeExpr :: Expr -> ReadWriteList +analyzeExpr = \case + Var x -> readVariable x + Lit _ -> mempty + UnOp _ e -> analyzeExpr e + BinOp _ e1 e2 -> analyzeExpr e1 <> analyzeExpr e2 + Cond e1 e2 e3 -> analyzeExpr e1 <> analyzeExpr e2 <> analyzeExpr e3 + Lam args _ body -> + let ReadWriteList rs ws = analyzeStatements body + args' = S.fromList (map snd args) + in ReadWriteList (rs `S.difference` args') (ws `S.difference` args') + Call _ args -> mconcat (map analyzeExpr args) + CallExpr f args -> mconcat (map analyzeExpr (f : args)) + +analyzeLeftExpr :: LeftExpr -> ReadWriteList +analyzeLeftExpr = \case + LeftVar x -> writeVariable x + LeftAt e1 e2 -> analyzeLeftExpr e1 <> analyzeExpr e2 + LeftGet _ e -> analyzeLeftExpr e + +analyzeAssignExpr :: AssignExpr -> ReadWriteList +analyzeAssignExpr = \case + AssignExpr _ e1 e2 -> analyzeLeftExpr e1 <> analyzeExpr e2 + AssignIncr e -> analyzeLeftExpr e + AssignDecr e -> analyzeLeftExpr e + +analyzeStatement :: Statement -> ReadWriteList +analyzeStatement = \case + ExprStatement e -> analyzeExpr e + Block body -> analyzeStatements body + If e body1 body2 -> analyzeExpr e <> analyzeStatements body1 <> maybe mempty analyzeStatements body2 + For _ x init pred incr body -> writeVariable x <> analyzeExpr init <> analyzeExpr pred <> analyzeAssignExpr incr <> analyzeStatements body + ForEach _ x e body -> writeVariable x <> analyzeExpr e <> analyzeStatements body + While e body -> analyzeExpr e <> analyzeStatements body + Declare _ x init -> + writeVariable x <> case init of + DeclareDefault -> mempty + DeclareCopy e -> analyzeExpr e + DeclareInitialize es -> mconcat (map analyzeExpr es) + DeclareDestructure xs e -> mconcat (map writeVariable xs) <> analyzeExpr e + Assign e -> analyzeAssignExpr e + Assert e -> analyzeExpr e + Return e -> analyzeExpr e + +analyzeStatements :: [Statement] -> ReadWriteList +analyzeStatements = mconcat . map analyzeStatement diff --git a/src/Jikka/Common/Alpha.hs b/src/Jikka/Common/Alpha.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/Alpha.hs @@ -0,0 +1,86 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : Jikka.Common.Alpha +-- Description : provides a monad to run alpha-conversion. / alpha 変換用のモナドを提供します。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- `Jikka.Common.Alpha` provides a monad to run alpha-conversion. This monad has only a feature to make unique numbers. +module Jikka.Common.Alpha where + +import Control.Arrow (first) +import Control.Monad.Except +import Control.Monad.Identity (Identity (..)) +import Control.Monad.Reader +import Control.Monad.Signatures +import Control.Monad.State.Strict +import Control.Monad.Writer.Strict + +class Monad m => MonadAlpha m where + nextCounter :: m Int + +newtype AlphaT m a = AlphaT {runAlphaT :: Int -> m (a, Int)} + +instance Monad m => MonadAlpha (AlphaT m) where + nextCounter = AlphaT (\i -> return (i, i + 1)) + +instance Functor m => Functor (AlphaT m) where + fmap f (AlphaT x) = AlphaT (\i -> fmap (first f) (x i)) + +instance Monad m => Applicative (AlphaT m) where + pure x = AlphaT (\i -> pure (x, i)) + AlphaT f <*> AlphaT x = AlphaT $ \i -> do + (f, i) <- f i + (x, i) <- x i + return (f x, i) + +instance Monad m => Monad (AlphaT m) where + AlphaT x >>= f = AlphaT $ \i -> do + (x, i) <- x i + runAlphaT (f x) i + +instance MonadFix m => MonadFix (AlphaT m) where + mfix f = AlphaT (\i -> mfix (\x -> runAlphaT (f (fst x)) i)) + +liftCatch :: Catch e m (a, Int) -> Catch e (AlphaT m) a +liftCatch catchE m h = AlphaT (\i -> runAlphaT m i `catchE` \e -> runAlphaT (h e) i) + +instance MonadTrans AlphaT where + lift m = AlphaT $ \i -> do + a <- m + return (a, i) + +instance MonadError e m => MonadError e (AlphaT m) where + throwError = lift . throwError + catchError = liftCatch catchError + +instance MonadIO m => MonadIO (AlphaT m) where + liftIO = lift . liftIO + +evalAlphaT :: Functor m => AlphaT m a -> Int -> m a +evalAlphaT f i = fst <$> runAlphaT f i + +instance MonadAlpha m => MonadAlpha (ExceptT e m) where + nextCounter = lift nextCounter + +instance MonadAlpha m => MonadAlpha (ReaderT r m) where + nextCounter = lift nextCounter + +instance MonadAlpha m => MonadAlpha (StateT s m) where + nextCounter = lift nextCounter + +instance (MonadAlpha m, Monoid w) => MonadAlpha (WriterT w m) where + nextCounter = lift nextCounter + +evalAlpha :: AlphaT Identity a -> Int -> a +evalAlpha f i = runIdentity (evalAlphaT f i) + +resetAlphaT :: Monad m => Int -> AlphaT m () +resetAlphaT i = AlphaT $ \_ -> return ((), i) diff --git a/src/Jikka/Common/Combinatorics.hs b/src/Jikka/Common/Combinatorics.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/Combinatorics.hs @@ -0,0 +1,18 @@ +module Jikka.Common.Combinatorics where + +fact :: Integral a => a -> a +fact n | n < 0 = error "Jikka.Common.Combinatorics.fact: invalid argument" +fact n = product [1 .. n] + +choose :: Integral a => a -> a -> a +choose n r | not (0 <= r && r <= n) = error "Jikka.Common.Combinatorics.choose: invalid argument" +choose n r = product [n - r + 1 .. n] `div` product [1 .. r] + +permute :: Integral a => a -> a -> a +permute n r | not (0 <= r && r <= n) = error "Jikka.Common.Combinatorics.permute: invalid argument" +permute n r = product [n - r + 1 .. n] + +multichoose :: Integral a => a -> a -> a +multichoose n r | not (0 <= r && r <= n) = error "Jikka.Common.Combinatorics.multichoose: invalid argument" +multichoose 0 0 = 1 +multichoose n r = choose (n + r - 1) r diff --git a/src/Jikka/Common/Error.hs b/src/Jikka/Common/Error.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/Error.hs @@ -0,0 +1,296 @@ +{-# LANGUAGE FlexibleContexts #-} + +-- | +-- Module : Jikka.Common.Error +-- Description : provides a data type which represents various errors. / 種々のエラーを表現するデータ型を提供します。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Common.Error + ( module Control.Monad.Except, + + -- * error data types + Responsibility (..), + ErrorGroup (..), + Error (..), + + -- * general utilities for `Control.Monad.Except` + wrapError, + wrapError', + wrapAt, + wrapAt', + maybeToError, + eitherToError, + + -- * utilities to report multiple errors + catchError', + reportErrors, + reportErrors2, + reportErrors3, + reportErrors4, + reportErrors5, + + -- * function to construct errors + lexicalError, + lexicalErrorAt, + syntaxError, + syntaxErrorAt, + symbolError, + symbolErrorAt, + typeError, + semanticError, + evaluationError, + runtimeError, + assertionError, + commandLineError, + wrongInputError, + internalError, + + -- * actions to throw errors + throwLexicalError, + throwLexicalErrorAt, + throwSyntaxError, + throwSyntaxErrorAt, + throwSyntaxErrorAt', + throwSymbolError, + throwSymbolErrorAt, + throwSymbolErrorAt', + throwTypeError, + throwTypeErrorAt, + throwTypeErrorAt', + throwSemanticError, + throwSemanticErrorAt, + throwSemanticErrorAt', + throwEvaluationError, + throwRuntimeError, + throwRuntimeErrorAt, + throwRuntimeErrorAt', + throwAssertionError, + throwCommandLineError, + throwWrongInputError, + throwInternalError, + throwInternalErrorAt, + throwInternalErrorAt', + + -- * utilities for other types of errors + bug, + todo, + ) +where + +import Control.Monad.Except +import Data.Either (isRight, lefts, rights) +import Jikka.Common.Location + +data Responsibility + = UserMistake + | ImplementationBug + deriving (Eq, Ord, Show, Read) + +data ErrorGroup + = -- | It's impossible to split the given source text into tokens. + LexicalError + | -- | It's impossible to construct AST from tokens. + SyntaxError + | -- | There are undefined variables or functions in AST. + SymbolError + | -- | It's impossible reconstruct types for AST. + TypeError + | -- | other semantic erros + SemanticError + | -- | User's program are not ready to evaluate. + EvaluationError + | -- | User's program failed while running. + RuntimeError + | -- | User's program violates its assertion. + AssertionError + | -- | The given command line arguments are not acceptable. + CommandLineError + | -- | User's program was correctly running but wrong input text is given. + WrongInputError + | -- | It's an bug of implementation. + InternalError + deriving (Eq, Ord, Show, Read) + +data Error + = Error String + | ErrorAppend Error Error + | WithGroup ErrorGroup Error + | WithWrapped String Error + | WithLocation Loc Error + | WithResponsibility Responsibility Error + deriving (Eq, Ord, Show, Read) + +instance Semigroup Error where + (<>) = ErrorAppend + +-- | The list must be non-empty. +errorList :: [Error] -> Error +errorList [] = bug "The list must be non-empty." +errorList (err : errs) = foldl ErrorAppend err errs + +wrapError :: MonadError e m => (e -> e) -> m a -> m a +wrapError wrap f = f `catchError` (\err -> throwError (wrap err)) + +wrapError' :: MonadError Error m => String -> m a -> m a +wrapError' message f = wrapError (WithWrapped message) f + +wrapAt :: MonadError Error m => Loc -> m a -> m a +wrapAt loc = wrapError (WithLocation loc) + +wrapAt' :: MonadError Error m => Maybe Loc -> m a -> m a +wrapAt' loc = maybe id wrapAt loc + +maybeToError :: MonadError a m => a -> Maybe b -> m b +maybeToError a Nothing = throwError a +maybeToError _ (Just b) = return b + +eitherToError :: MonadError a m => Either a b -> m b +eitherToError = liftEither + +-- | `catchError'` is the inverse of `liftError`. +catchError' :: MonadError e m => m a -> m (Either e a) +catchError' f = (Right <$> f) `catchError` (\err -> return (Left err)) + +reportErrors :: MonadError Error m => [Either Error a] -> m [a] +reportErrors xs + | all isRight xs = return $ rights xs + | otherwise = throwError $ errorList (lefts xs) + +reportErrors2 :: MonadError Error m => Either Error a -> Either Error b -> m (a, b) +reportErrors2 (Right a) (Right b) = return (a, b) +reportErrors2 a b = throwError $ errorList (lefts [() <$ a, () <$ b]) + +reportErrors3 :: MonadError Error m => Either Error a -> Either Error b -> Either Error c -> m (a, b, c) +reportErrors3 (Right a) (Right b) (Right c) = return (a, b, c) +reportErrors3 a b c = throwError $ errorList (lefts [() <$ a, () <$ b, () <$ c]) + +reportErrors4 :: MonadError Error m => Either Error a -> Either Error b -> Either Error c -> Either Error d -> m (a, b, c, d) +reportErrors4 (Right a) (Right b) (Right c) (Right d) = return (a, b, c, d) +reportErrors4 a b c d = throwError $ errorList (lefts [() <$ a, () <$ b, () <$ c, () <$ d]) + +reportErrors5 :: MonadError Error m => Either Error a -> Either Error b -> Either Error c -> Either Error d -> Either Error e -> m (a, b, c, d, e) +reportErrors5 (Right a) (Right b) (Right c) (Right d) (Right e) = return (a, b, c, d, e) +reportErrors5 a b c d e = throwError $ errorList (lefts [() <$ a, () <$ b, () <$ c, () <$ d, () <$ e]) + +lexicalError :: String -> Error +lexicalError = WithGroup LexicalError . Error + +lexicalErrorAt :: Loc -> String -> Error +lexicalErrorAt loc = WithLocation loc . WithGroup LexicalError . Error + +syntaxError :: String -> Error +syntaxError = WithGroup SyntaxError . Error + +syntaxErrorAt :: Loc -> String -> Error +syntaxErrorAt loc = WithLocation loc . WithGroup SyntaxError . Error + +symbolError :: String -> Error +symbolError = WithGroup SymbolError . Error + +symbolErrorAt :: Loc -> String -> Error +symbolErrorAt loc = WithLocation loc . WithGroup SymbolError . Error + +typeError :: String -> Error +typeError = WithGroup TypeError . Error + +semanticError :: String -> Error +semanticError = WithGroup SemanticError . Error + +evaluationError :: String -> Error +evaluationError = WithGroup EvaluationError . Error + +runtimeError :: String -> Error +runtimeError = WithGroup RuntimeError . Error + +assertionError :: String -> Error +assertionError = WithGroup AssertionError . Error + +commandLineError :: String -> Error +commandLineError = WithGroup CommandLineError . Error + +wrongInputError :: String -> Error +wrongInputError = WithGroup WrongInputError . Error + +internalError :: String -> Error +internalError = WithGroup InternalError . Error + +throwLexicalError :: MonadError Error m => String -> m a +throwLexicalError = throwError . WithGroup LexicalError . Error + +throwLexicalErrorAt :: MonadError Error m => Loc -> String -> m a +throwLexicalErrorAt loc = throwError . WithLocation loc . WithGroup LexicalError . Error + +throwSyntaxError :: MonadError Error m => String -> m a +throwSyntaxError = throwError . WithGroup SyntaxError . Error + +throwSyntaxErrorAt :: MonadError Error m => Loc -> String -> m a +throwSyntaxErrorAt loc = throwError . WithLocation loc . WithGroup SyntaxError . Error + +throwSyntaxErrorAt' :: MonadError Error m => Maybe Loc -> String -> m a +throwSyntaxErrorAt' loc = throwError . maybe id WithLocation loc . WithGroup SyntaxError . Error + +throwSymbolError :: MonadError Error m => String -> m a +throwSymbolError = throwError . WithGroup SymbolError . Error + +throwSymbolErrorAt :: MonadError Error m => Loc -> String -> m a +throwSymbolErrorAt loc = throwError . WithLocation loc . WithGroup SymbolError . Error + +throwSymbolErrorAt' :: MonadError Error m => Maybe Loc -> String -> m a +throwSymbolErrorAt' loc = throwError . maybe id WithLocation loc . WithGroup SymbolError . Error + +throwTypeError :: MonadError Error m => String -> m a +throwTypeError = throwError . WithGroup TypeError . Error + +throwTypeErrorAt :: MonadError Error m => Loc -> String -> m a +throwTypeErrorAt loc = throwError . WithLocation loc . WithGroup TypeError . Error + +throwTypeErrorAt' :: MonadError Error m => Maybe Loc -> String -> m a +throwTypeErrorAt' loc = throwError . maybe id WithLocation loc . WithGroup TypeError . Error + +throwSemanticError :: MonadError Error m => String -> m a +throwSemanticError = throwError . WithGroup SemanticError . Error + +throwSemanticErrorAt :: MonadError Error m => Loc -> String -> m a +throwSemanticErrorAt loc = throwError . WithLocation loc . WithGroup SemanticError . Error + +throwSemanticErrorAt' :: MonadError Error m => Maybe Loc -> String -> m a +throwSemanticErrorAt' loc = throwError . maybe id WithLocation loc . WithGroup SemanticError . Error + +throwEvaluationError :: MonadError Error m => String -> m a +throwEvaluationError = throwError . WithGroup EvaluationError . Error + +throwRuntimeError :: MonadError Error m => String -> m a +throwRuntimeError = throwError . WithGroup RuntimeError . Error + +throwRuntimeErrorAt :: MonadError Error m => Loc -> String -> m a +throwRuntimeErrorAt loc = throwError . WithLocation loc . WithGroup RuntimeError . Error + +throwRuntimeErrorAt' :: MonadError Error m => Maybe Loc -> String -> m a +throwRuntimeErrorAt' loc = throwError . maybe id WithLocation loc . WithGroup RuntimeError . Error + +throwAssertionError :: MonadError Error m => String -> m a +throwAssertionError = throwError . WithGroup AssertionError . Error + +throwCommandLineError :: MonadError Error m => String -> m a +throwCommandLineError = throwError . WithGroup CommandLineError . Error + +throwWrongInputError :: MonadError Error m => String -> m a +throwWrongInputError = throwError . WithGroup WrongInputError . Error + +throwInternalError :: MonadError Error m => String -> m a +throwInternalError = throwError . WithGroup InternalError . Error + +throwInternalErrorAt :: MonadError Error m => Loc -> String -> m a +throwInternalErrorAt loc = throwError . WithLocation loc . WithGroup InternalError . Error + +throwInternalErrorAt' :: MonadError Error m => Maybe Loc -> String -> m a +throwInternalErrorAt' loc = throwError . maybe id WithLocation loc . WithGroup InternalError . Error + +bug :: String -> a +bug msg = error $ "Fatal Error (implementation's bug): " ++ msg + +todo :: String -> a +todo msg = error $ "TODO Error (the feature is not implemented yet): " ++ msg diff --git a/src/Jikka/Common/Format/AutoIndent.hs b/src/Jikka/Common/Format/AutoIndent.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/Format/AutoIndent.hs @@ -0,0 +1,31 @@ +module Jikka.Common.Format.AutoIndent where + +import Data.List (isPrefixOf, isSuffixOf) + +indent :: String +indent = "" + +dedent :: String +dedent = "" + +makeIndentFromMarkers :: Int -> [String] -> [String] +makeIndentFromMarkers size = go 0 + where + go :: Int -> [String] -> [String] + go _ [] = [] + go n (line : lines) + | line == indent = go (n + size) lines + | line == dedent = go (n - size) lines + | otherwise = (replicate n ' ' ++ line) : go n lines + +makeIndentFromBraces :: Int -> [String] -> [String] +makeIndentFromBraces size = makeIndentFromMarkers size . insertIndentDedentFromBraces + +insertIndentDedentFromBraces :: [String] -> [String] +insertIndentDedentFromBraces = concatMap go + where + go :: String -> [String] + go line = + let close = if "}" `isPrefixOf` line then [dedent] else [] + open = if "{" `isSuffixOf` line then [indent] else [] + in close ++ [line] ++ open diff --git a/src/Jikka/Common/Format/Color.hs b/src/Jikka/Common/Format/Color.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/Format/Color.hs @@ -0,0 +1,29 @@ +module Jikka.Common.Format.Color + ( ColorFlag (..), + withColor, + withBold, + hGetColorFlag, + Color (..), + ) +where + +import System.Console.ANSI +import System.IO (Handle) + +data ColorFlag + = EnableColor + | DisableColor + deriving (Eq, Ord, Show, Read) + +withColor :: ColorFlag -> Color -> String -> String +withColor DisableColor _ s = s +withColor EnableColor color s = setSGRCode [SetColor Foreground Vivid color] ++ s ++ setSGRCode [SetColor Foreground Dull White] + +withBold :: ColorFlag -> String -> String +withBold DisableColor s = s +withBold EnableColor s = setSGRCode [SetConsoleIntensity BoldIntensity] ++ s ++ setSGRCode [SetConsoleIntensity NormalIntensity] + +hGetColorFlag :: Handle -> IO ColorFlag +hGetColorFlag handle = do + supported <- hSupportsANSI handle + return (if supported then EnableColor else DisableColor) diff --git a/src/Jikka/Common/Format/Error.hs b/src/Jikka/Common/Format/Error.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/Format/Error.hs @@ -0,0 +1,135 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +module Jikka.Common.Format.Error + ( prettyError, + prettyError', + prettyErrorWithText, + hPrintError, + hPrintErrorWithText, + ) +where + +import Data.List (intercalate) +import Data.Maybe (fromMaybe) +import Data.Text (Text) +import Jikka.Common.Error +import Jikka.Common.Format.Color +import Jikka.Common.Format.Location +import Jikka.Common.Location +import System.IO (Handle, hPutStrLn) + +-- | `unpackCombinedErrors` removes `ErrorAppend` ctor from the given `Error`. +unpackCombinedErrors :: Error -> [Error] +unpackCombinedErrors = go + where + go :: Error -> [Error] + go = \case + err@(Error _) -> [err] + ErrorAppend err1 err2 -> go err1 ++ go err2 + WithGroup group err -> map (WithGroup group) (go err) + WithWrapped msg err -> map (WithWrapped msg) (go err) + WithLocation loc err -> map (WithLocation loc) (go err) + WithResponsibility resp err -> map (WithResponsibility resp) (go err) + +prettyError :: ColorFlag -> Error -> [String] +prettyError color = map (prettyError1 color) . unpackCombinedErrors + +prettyError' :: Error -> [String] +prettyError' = prettyError DisableColor + +-- | @err@ must not have `ErrorAppend`. +prettyError1 :: ColorFlag -> Error -> String +prettyError1 color err = intercalate ": " ((group ++ loc ++ resp) : getMessages err) + where + group = withColor color Red $ prettyGroup (getErrorGroup err) + loc = case getLocation err of + Nothing -> "" + Just loc -> " (" ++ prettyLoc loc ++ ")" + resp = case getResponsibility err of + Just UserMistake -> " (user's mistake?)" + Just ImplementationBug -> " (implementation's bug?)" + Nothing -> "" + +prettyErrorWithText :: ColorFlag -> Text -> Error -> [String] +prettyErrorWithText color text = intercalate [""] . map (prettyErrorWithText1 color text) . unpackCombinedErrors + +-- | @err@ must not have `ErrorAppend`. +prettyErrorWithText1 :: ColorFlag -> Text -> Error -> [String] +prettyErrorWithText1 color text err = case getLocation err of + Nothing -> [prettyError1 color err] + Just loc -> prettyError1 color err : prettyLocWithText color text loc + +prettyGroup :: Maybe ErrorGroup -> String +prettyGroup = \case + Nothing -> "Error" + Just LexicalError -> "Lexical Error" + Just SyntaxError -> "Syntax Error" + Just SemanticError -> "Semantic Error" + Just SymbolError -> "Symbol Error" + Just TypeError -> "Type Error" + Just EvaluationError -> "Evaluation Error" + Just RuntimeError -> "Runtime Error" + Just AssertionError -> "Assertion Error" + Just CommandLineError -> "Command Line Error" + Just WrongInputError -> "Wrong Input Error" + Just InternalError -> "Internal Error" + +-- | @err@ must not have `ErrorAppend`. +getMessages :: Error -> [String] +getMessages = \case + Error message -> [message] + ErrorAppend _ _ -> bug "ErrorAppend is not allowed here." + WithGroup _ err -> getMessages err + WithWrapped message err -> message : getMessages err + WithLocation _ err -> getMessages err + WithResponsibility _ err -> getMessages err + +-- | @err@ must not have `ErrorAppend`. +getErrorGroup :: Error -> Maybe ErrorGroup +getErrorGroup = \case + Error _ -> Nothing + ErrorAppend _ _ -> bug "ErrorAppend is not allowed here." + WithGroup group err -> Just (fromMaybe group (getErrorGroup err)) + WithWrapped _ err -> getErrorGroup err + WithLocation _ err -> getErrorGroup err + WithResponsibility _ err -> getErrorGroup err + +-- | @err@ must not have `ErrorAppend`. +getLocation :: Error -> Maybe Loc +getLocation = \case + Error _ -> Nothing + ErrorAppend _ _ -> bug "ErrorAppend is not allowed here." + WithGroup _ err -> getLocation err + WithWrapped _ err -> getLocation err + WithLocation loc err -> Just (fromMaybe loc (getLocation err)) + WithResponsibility _ err -> getLocation err + +getResponsibilityFromErrorGroup :: ErrorGroup -> Maybe Responsibility +getResponsibilityFromErrorGroup = \case + CommandLineError -> Nothing + WrongInputError -> Nothing + InternalError -> Just ImplementationBug + _ -> Just UserMistake + +-- | @err@ must not have `ErrorAppend`. +getResponsibility :: Error -> Maybe Responsibility +getResponsibility = \case + Error _ -> Nothing + ErrorAppend _ _ -> bug "ErrorAppend is not allowed here." + WithGroup group err -> case getResponsibility err of + Just resp -> Just resp + Nothing -> getResponsibilityFromErrorGroup group + WithWrapped _ err -> getResponsibility err + WithLocation _ err -> getResponsibility err + WithResponsibility resp _ -> Just resp + +hPrintError :: Handle -> Error -> IO () +hPrintError handle err = do + color <- hGetColorFlag handle + mapM_ (hPutStrLn handle) (prettyError color err) + +hPrintErrorWithText :: Handle -> Text -> Error -> IO () +hPrintErrorWithText handle text err = do + color <- hGetColorFlag handle + mapM_ (hPutStrLn handle) (prettyErrorWithText color text err) diff --git a/src/Jikka/Common/Format/Location.hs b/src/Jikka/Common/Format/Location.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/Format/Location.hs @@ -0,0 +1,37 @@ +module Jikka.Common.Format.Location + ( prettyLoc, + prettyLocWithText, + ) +where + +import Data.Text (Text) +import qualified Data.Text as T +import Jikka.Common.Format.Color +import Jikka.Common.Location + +prettyLoc :: Loc -> String +prettyLoc loc = "line " ++ show (line loc) ++ " column " ++ show (column loc) + +prettyLocWithText :: ColorFlag -> Text -> Loc -> [String] +prettyLocWithText color text (Loc y x width) = result + where + lines :: [Text] + lines = T.lines text + paddingSize :: Int + paddingSize = length $ show (y + 1) + padRight :: String -> String + padRight s = s ++ replicate (paddingSize - length s) ' ' + prettyLine :: (String -> String) -> Int -> [String] + prettyLine f y + | 1 <= y && y <= length lines = [withColor color Blue (padRight (show y) ++ " |") ++ f (T.unpack (lines !! (y - 1)))] + | otherwise = [] + result :: [String] + result + | 1 <= y && y <= length lines = + concat + [ prettyLine id (y - 1), + prettyLine (\line -> take (x - 1) line ++ withColor color Red (take width (drop (x - 1) line)) ++ drop (x - 1 + width) line) y, + [replicate (paddingSize + 2 + x - 1) ' ' ++ withColor color Red (replicate (max 1 width) '^')], + prettyLine id (y + 1) + ] + | otherwise = [""] diff --git a/src/Jikka/Common/Format/Show.hs b/src/Jikka/Common/Format/Show.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/Format/Show.hs @@ -0,0 +1,6 @@ +module Jikka.Common.Format.Show where + +import Data.Text (Text, pack) + +run :: Show program => program -> Either String Text +run e = Right . pack $ show e diff --git a/src/Jikka/Common/Graph.hs b/src/Jikka/Common/Graph.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/Graph.hs @@ -0,0 +1,81 @@ +module Jikka.Common.Graph where + +import Control.Monad +import Control.Monad.ST +import Data.List (nub) +import Data.STRef +import qualified Data.Vector as V +import qualified Data.Vector.Mutable as MV + +type Graph = V.Vector [Int] + +makeReversedDigraph :: Graph -> Graph +makeReversedDigraph g = runST $ do + let n = V.length g + h <- MV.replicate n [] + forM_ [0 .. n - 1] $ \x -> do + forM_ (g V.! x) $ \y -> do + MV.modify h (x :) y + V.freeze h + +makeInducedDigraph :: Graph -> V.Vector Int -> Graph +makeInducedDigraph g f = runST $ do + let n = V.length g + let k = if V.null f then 0 else V.maximum f + 1 + h <- MV.replicate k [] + forM_ [0 .. n - 1] $ \x -> do + forM_ (g V.! x) $ \y -> do + MV.modify h ((f V.! y) :) (f V.! x) + forM_ [0 .. k - 1] $ \a -> do + MV.modify h nub a + V.freeze h + +-- | `decomposeToStronglyConnectedComponents` does SCC in \(O(V + E)\) using Kosaraju's algorithm. +-- It takes a digraph \(G = (V, E)\) as an adjacent list \(g : V \to V^{\lt \omega}\), and returns an mapping \(f : V \to V'\) for the SCC DAG \(G' = (V', E')\). +-- The indices of vertices of the SCC DAG are topologically sorted. +decomposeToStronglyConnectedComponents :: Graph -> V.Vector Int +decomposeToStronglyConnectedComponents g = runST $ do + let n = V.length g + let unless' used x f = do + usedX <- MV.read used x + unless usedX $ do + f + -- The first DFS + let order = topologicalSort g + -- DFS on the reversed graph + let gRev = makeReversedDigraph g + componentOf <- MV.replicate n (-1) + size <- newSTRef 0 + used <- MV.replicate n False + let go x = do + MV.write used x True + forM_ (gRev V.! x) $ \y -> do + unless' used y $ do + go y + V.forM_ order $ \x -> do + unless' used x $ do + go x + modifySTRef' size succ + V.freeze componentOf + +-- | `topologicalSort` does topological sort in \(O(V + E)\) using Tarjan's algorithm. +-- The input is an adjacent list of a DAG. +topologicalSort :: Graph -> V.Vector Int +topologicalSort g = runST $ do + let n = V.length g + let unless' used x f = do + usedX <- MV.read used x + unless usedX $ do + f + order <- newSTRef [] + used <- MV.replicate n False + let go x = do + MV.write used x True + forM_ (g V.! x) $ \y -> do + unless' used y $ do + go y + modifySTRef' order (x :) + forM_ [0 .. n - 1] $ \x -> do + unless' used x $ do + go x + V.fromList <$> readSTRef order diff --git a/src/Jikka/Common/IO.hs b/src/Jikka/Common/IO.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/IO.hs @@ -0,0 +1,27 @@ +module Jikka.Common.IO where + +import Data.Char (isSpace) +import System.IO + +hTakeWhile :: Handle -> (Char -> Bool) -> IO String +hTakeWhile handle pred = do + isEOF <- hIsEOF handle + if isEOF + then return "" + else do + c <- hLookAhead handle + if pred c + then do + _ <- hGetChar handle + (c :) <$> hTakeWhile handle pred + else return "" + +hGetWord :: Handle -> IO String +hGetWord handle = do + hTakeWhile handle isSpace + c <- hGetChar handle + s <- hTakeWhile handle (not . isSpace) + return (c : s) + +getWord :: IO String +getWord = hGetWord stdin diff --git a/src/Jikka/Common/IOFormat.hs b/src/Jikka/Common/IOFormat.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/IOFormat.hs @@ -0,0 +1,221 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +module Jikka.Common.IOFormat where + +import Control.Arrow +import Control.Monad.Identity +import Data.IORef +import Data.List +import qualified Data.Map as M +import qualified Data.Vector as V +import Jikka.Common.Error +import Jikka.Common.IO (hGetWord) +import System.IO (stdin) +import Text.Read (readMaybe) + +data FormatExpr + = Var String + | Plus FormatExpr Integer + | At FormatExpr String + | Len FormatExpr + deriving (Eq, Ord, Read, Show) + +data FormatTree + = Exp FormatExpr + | Newline + | Seq [FormatTree] + | Loop String FormatExpr FormatTree + deriving (Eq, Ord, Read, Show) + +data IOFormat = IOFormat + { inputVariables :: [String], + inputTree :: FormatTree, + -- | This uses `Either` to distinguish a type and the 1-tuple of it. + outputVariables :: Either String [String], + outputTree :: FormatTree + } + deriving (Eq, Ord, Read, Show) + +mapFormatTreeM :: Monad m => (FormatTree -> m FormatTree) -> FormatTree -> m FormatTree +mapFormatTreeM f = \case + Loop i n body -> do + body <- mapFormatTreeM f body + f $ Loop i n body + Seq formats -> Seq <$> mapM f formats + format -> f format + +mapFormatTree :: (FormatTree -> FormatTree) -> FormatTree -> FormatTree +mapFormatTree f = runIdentity . mapFormatTreeM (return . f) + +normalizeFormatTree :: FormatTree -> FormatTree +normalizeFormatTree = \case + Exp e -> Exp e + Newline -> Newline + Seq formats -> + let unSeq = \case + Seq formats -> formats + format -> [format] + in Seq (concatMap (unSeq . normalizeFormatTree) formats) + Loop i n body -> case normalizeFormatTree body of + Seq [] -> Seq [] + body -> Loop i n body + +normalizeIOFormat :: IOFormat -> IOFormat +normalizeIOFormat format = + format + { inputTree = normalizeFormatTree (inputTree format), + outputTree = normalizeFormatTree (outputTree format) + } + +hasNewline :: FormatTree -> Bool +hasNewline = \case + Exp _ -> False + Newline -> True + Seq formats -> any hasNewline formats + Loop _ _ body -> hasNewline body + +formatFormatExpr :: FormatExpr -> String +formatFormatExpr = \case + Var x -> x + Plus e k -> "(" ++ formatFormatExpr e ++ " + " ++ show k ++ ")" + At e i -> formatFormatExpr e ++ "[" ++ i ++ "]" + Len e -> "len(" ++ formatFormatExpr e ++ ")" + +formatFormatTree :: FormatTree -> String +formatFormatTree = + let replace :: Eq a => [a] -> [a] -> [a] -> [a] + replace patt subst = go + where + go text | patt `isPrefixOf` text = subst ++ go (drop (length patt) text) + go [] = [] + go (c : s) = c : go s + unwords' = replace "\n\n" "\n" . replace "\n " "\n" . replace " \n" "\n" . unwords + in \case + Exp e -> formatFormatExpr e + Newline -> "(newline)\n" + Seq formats -> unwords' (map formatFormatTree formats) + Loop i n body -> + unwords' + [ "for " ++ i ++ " < " ++ formatFormatExpr n ++ " {\n", + formatFormatTree body ++ "\n", + "}" + ] + +formatIOFormat :: IOFormat -> String +formatIOFormat format = + unlines + ( [ "input tree:" + ] + ++ map (" " ++) (lines (formatFormatTree (inputTree format))) + ++ [ "input variables: " ++ show (inputVariables format), + "output variables: " ++ show (outputVariables format), + "output tree:" + ] + ++ map (" " ++) (lines (formatFormatTree (outputTree format))) + ) + +packSubscriptedVar :: String -> [String] -> FormatExpr +packSubscriptedVar x indices = foldl At (Var x) indices + +packSubscriptedVar' :: String -> [String] -> FormatTree +packSubscriptedVar' = (Exp .) . packSubscriptedVar + +unpackSubscriptedVar :: MonadError Error m => FormatExpr -> m (String, [String]) +unpackSubscriptedVar = \case + Var x -> return (x, []) + At e i -> second (++ [i]) <$> unpackSubscriptedVar e + e -> throwInternalError $ "not a subscripted variable: " ++ formatFormatExpr e + +makeReadValueIO :: (MonadError Error m, MonadIO m) => (value -> m Integer) -> (Integer -> value) -> (value -> m (V.Vector value)) -> (V.Vector value -> value) -> IOFormat -> m ([value], M.Map String value) +makeReadValueIO toInt fromInt toList fromList format = wrapError' "Jikka.Common.IOFormat.makeReadValueIO" $ do + env <- liftIO $ newIORef M.empty + sizes <- liftIO $ newIORef M.empty + let lookup x = do + y <- M.lookup x <$> liftIO (readIORef env) + case y of + Nothing -> throwInternalError $ "undefined variable: " ++ x + Just y -> return y + let go = \case + Exp e -> do + (x, indices) <- unpackSubscriptedVar e + word <- liftIO $ hGetWord stdin + n <- case readMaybe word of + Just n -> return n + Nothing -> throwWrongInputError $ "not a integer: " ++ word + y <- M.lookup x <$> liftIO (readIORef env) + y <- case y of + Just y -> return y + Nothing -> do + let go' x i = do + size <- M.lookup i <$> liftIO (readIORef sizes) + case size of + Nothing -> throwInternalError $ "undefined variable: " ++ i + Just size -> return . fromList $ V.replicate (fromInteger size) x + foldM go' (fromInt (-1)) indices + let go' y = \case + [] -> return (fromInt n) + (i : indices) -> do + i <- toInt =<< lookup i + y <- toList y + z <- go' (y V.! fromInteger i) indices + return . fromList $ y V.// [(fromInteger i, z)] + y <- go' y indices + liftIO $ modifyIORef' env (M.insert x y) + Newline -> return () + Seq formats -> mapM_ go formats + Loop i n body -> do + n <- case n of + Var n -> toInt =<< lookup n + Plus (Var n) k -> (+ k) <$> (toInt =<< lookup n) + Len (Var xs) -> toInteger . V.length <$> (toList =<< lookup xs) + _ -> throwInternalError $ "invalid loop size in input tree: " ++ formatFormatExpr n + liftIO $ modifyIORef' sizes (M.insert i n) + forM_ [0 .. n -1] $ \i' -> do + liftIO $ modifyIORef' env (M.insert i (fromInt i')) + go body + go (inputTree format) + values <- mapM lookup (inputVariables format) + env <- liftIO $ readIORef env + return (values, env) + +makeWriteValueIO :: (MonadError Error m, MonadIO m) => (value -> m [value]) -> (Integer -> value) -> (value -> m Integer) -> (value -> m (V.Vector value)) -> IOFormat -> M.Map String value -> value -> m () +makeWriteValueIO toTuple fromInt toInt toList format env value = wrapError' "Jikka.Common.IOFormat.makeWriteValueIO" $ do + env <- liftIO $ newIORef env + let lookup x = do + y <- M.lookup x <$> liftIO (readIORef env) + case y of + Nothing -> throwInternalError $ "undefined variable: " ++ x + Just y -> return y + case outputVariables format of + Left x -> liftIO $ modifyIORef' env (M.insert x value) + Right xs -> do + values <- toTuple value + when (length values /= length xs) $ do + throwRuntimeError $ "sizes of values mismtach: expected = " ++ show (length xs) ++ ", actual = " ++ show (length values) + forM_ (zip xs values) $ \(x, value) -> do + liftIO $ modifyIORef' env (M.insert x value) + let evaluate = \case + Var n -> lookup n + Plus e k -> fromInt . (+ k) <$> (toInt =<< evaluate e) + Len e -> do + e <- toList =<< evaluate e + return . fromInt . toInteger $ V.length e + At e i -> do + xs <- toList =<< evaluate e + i <- toInt =<< lookup i + case xs V.!? fromInteger i of + Nothing -> throwRuntimeError $ "length of list is shorter than expected: expected > " ++ show i ++ ", actual = " ++ show (V.length xs) + Just x -> return x + let go = \case + Exp e -> do + x <- toInt =<< evaluate e + liftIO $ putStr (show x ++ " ") + Newline -> liftIO $ putChar '\n' + Seq formats -> mapM_ go formats + Loop i n body -> do + n <- toInt =<< evaluate n + forM_ [0 .. n -1] $ \i' -> do + liftIO $ modifyIORef' env (M.insert i (fromInt i')) + go body + go (outputTree format) diff --git a/src/Jikka/Common/Location.hs b/src/Jikka/Common/Location.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/Location.hs @@ -0,0 +1,31 @@ +{-# LANGUAGE DeriveFunctor #-} + +module Jikka.Common.Location where + +import Data.String (IsString (..)) + +-- | `Loc` represents a location of something in the users' source code. `line` and `column` is 1-based. +data Loc = Loc + { line :: !Int, + column :: !Int, + width :: !Int + } + deriving (Eq, Ord, Show, Read) + +data WithLoc a = WithLoc + { loc :: !Loc, + value :: !a + } + deriving (Eq, Ord, Show, Read, Functor) + +data WithLoc' a = WithLoc' + { loc' :: !(Maybe Loc), + value' :: !a + } + deriving (Eq, Ord, Show, Read, Functor) + +instance IsString a => IsString (WithLoc' a) where + fromString = WithLoc' Nothing . fromString + +withoutLoc :: a -> WithLoc' a +withoutLoc = WithLoc' Nothing diff --git a/src/Jikka/Common/Matrix.hs b/src/Jikka/Common/Matrix.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/Matrix.hs @@ -0,0 +1,113 @@ +{-# LANGUAGE DeriveFunctor #-} + +module Jikka.Common.Matrix + ( Matrix, + unMatrix, + makeMatrix, + makeMatrix', + matsize, + matsize', + matcheck, + matzero, + matone, + matadd, + matmul, + matap, + matscalar, + matpow, + ) +where + +import Control.Monad +import Control.Monad.ST +import qualified Data.Vector as V +import qualified Data.Vector.Mutable as MV + +-- | `Matrix` is data for matrices. +-- It is guaranteed that internal arrays are not jagged arrays. +newtype Matrix a = Matrix (V.Vector (V.Vector a)) + deriving (Eq, Ord, Show, Functor) + +unMatrix :: Matrix a -> V.Vector (V.Vector a) +unMatrix (Matrix a) = a + +-- | `matsize` computes the size of a matrix. +matsize :: Matrix a -> (Int, Int) +matsize (Matrix a) = matsize' a + +-- | `matsize'` computes the size of a matrix. +-- This assumes inputs are matrices (`matcheck`). +matsize' :: V.Vector (V.Vector a) -> (Int, Int) +matsize' a = + if V.null a + then (0, 0) + else (V.length a, V.length (a V.! 0)) + +-- | `matcheck` checks a given vector of vectors is a matrix. +-- That is, this returns `False` for jagged arrays. +matcheck :: V.Vector (V.Vector a) -> Bool +matcheck a = + let (_, w) = matsize' a + in all (\row -> V.length row == w) (V.toList a) + +makeMatrix :: V.Vector (V.Vector a) -> Maybe (Matrix a) +makeMatrix a = if matcheck a then Just (Matrix a) else Nothing + +makeMatrix' :: V.Vector (V.Vector a) -> Matrix a +makeMatrix' a = case makeMatrix a of + Nothing -> error "Jikka.Common.Matrix.makeMatrix': the input is not a matrix" + Just a -> a + +matzero :: Num a => Int -> Matrix a +matzero n = Matrix $ V.replicate n (V.replicate n 0) + +matone :: Num a => Int -> Matrix a +matone n = Matrix $ V.generate n (\y -> V.generate n (\x -> if y == x then 1 else 0)) + +-- | `matadd` calculates the addition \(A + B\) of two matrices \(A, B\). +-- This assumes sizes of inputs match. +matadd :: Num a => Matrix a -> Matrix a -> Matrix a +matadd (Matrix a) (Matrix b) = + let (h, w) = matsize' a + in Matrix $ V.generate h (\y -> V.generate w (\x -> (a V.! y V.! x) + (b V.! y V.! x))) + +-- | `matmul` calculates the multiplication \(A B\)of two matrices \(A, B\). +-- This assumes sizes of inputs match. +matmul :: Num a => Matrix a -> Matrix a -> Matrix a +matmul (Matrix a) (Matrix b) = runST $ do + let (h, n) = matsize' a + let (_, w) = matsize' b + c <- MV.replicateM h (MV.replicate w 0) + forM_ [0 .. h - 1] $ \y -> do + forM_ [0 .. n - 1] $ \z -> do + forM_ [0 .. w - 1] $ \x -> do + let delta = (a V.! y V.! z) * (b V.! z V.! x) + row <- MV.read c y + MV.modify row (+ delta) x + Matrix . V.fromList <$> MV.foldrM' (\row c' -> (: c') <$> V.freeze row) [] c + +-- | `matap` calculates the multiplication \(A x\) of a matrix \(A\) and a vector \(x\). +-- This assumes sizes of inputs match. +matap :: Num a => Matrix a -> V.Vector a -> V.Vector a +matap (Matrix a) b = runST $ do + let (h, w) = matsize' a + c <- MV.replicate h 0 + forM_ [0 .. h - 1] $ \y -> do + forM_ [0 .. w - 1] $ \x -> do + let delta = (a V.! y V.! x) * (b V.! x) + MV.modify c (+ delta) y + V.freeze c + +matscalar :: Num a => a -> Matrix a -> Matrix a +matscalar a (Matrix b) = Matrix $ V.map (V.map (a *)) b + +-- | `matpow` calculates the power \(A^k\) of a matrix \(A\) and a natural number \(k\). +-- This assumes inputs are square matrices. +-- This fails for \(k \lt 0\). +matpow :: (Show a, Num a) => Matrix a -> Integer -> Matrix a +matpow _ k | k < 0 = error "cannot calculate a negative power for a monoid" +matpow x k = go unit x k + where + unit = let (h, _) = matsize x in matone h + go y _ 0 = y + go y x k = go (if k `mod` 2 == 1 then matmul y x else y) (matmul x x) (k `div` 2) diff --git a/src/Jikka/Common/ModInt.hs b/src/Jikka/Common/ModInt.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/ModInt.hs @@ -0,0 +1,38 @@ +module Jikka.Common.ModInt + ( ModInt, + toModInt, + fromModInt, + moduloOfModInt, + ) +where + +import Data.Monoid + +data ModInt = ModInt Integer (Maybe Integer) + deriving (Eq, Ord, Read, Show) + +toModInt :: Integer -> Integer -> ModInt +toModInt _ m | m <= 0 = error $ "Jikka.Common.ModInt.toModInt: modulo must be positive, but m = " ++ show m +toModInt a m = ModInt (a `mod` m) (Just m) + +fromModInt :: ModInt -> Integer +fromModInt (ModInt a _) = a + +moduloOfModInt :: ModInt -> Maybe Integer +moduloOfModInt (ModInt _ m) = m + +instance Num ModInt where + ModInt _ (Just m1) + ModInt _ (Just m2) | m1 /= m2 = error $ "Jikka.Common.ModInt.(+): modulo must be the same, but m1 = " ++ show m1 ++ " and m2 = " ++ show m2 + ModInt a m1 + ModInt b m2 = case getFirst (First m1 <> First m2) of + Nothing -> ModInt (a + b) Nothing + Just m -> ModInt (let c = a + b in if c >= m then c - m else c) (Just m) + ModInt _ (Just m1) * ModInt _ (Just m2) | m1 /= m2 = error $ "Jikka.Common.ModInt.(*): modulo must be the same, but m1 = " ++ show m1 ++ " and m2 = " ++ show m2 + ModInt a m1 * ModInt b m2 = case getFirst (First m1 <> First m2) of + Nothing -> ModInt (a * b) Nothing + Just m -> ModInt ((a * b) `mod` m) (Just m) + abs = error "Jikka.Common.ModInt.fromInteger: cannot call abs for modint" + signum = error "Jikka.Common.ModInt.fromInteger: cannot signum for modint" + fromInteger a = ModInt a Nothing + negate (ModInt a m) = case m of + Nothing -> ModInt (- a) m + Just m -> ModInt (if a == 0 then 0 else m - a) (Just m) diff --git a/src/Jikka/Common/Parse/JoinLines.hs b/src/Jikka/Common/Parse/JoinLines.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/Parse/JoinLines.hs @@ -0,0 +1,33 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Jikka.Common.Parse.JoinLines + ( joinLinesWithParens, + removeEmptyLines, + ) +where + +import Jikka.Common.Error +import Jikka.Common.Location + +joinLinesWithParens :: forall m a. (MonadError Error m, Show a) => (a -> Bool) -> (a -> Bool) -> (a -> Bool) -> [WithLoc a] -> m [WithLoc a] +joinLinesWithParens isOpen isClose isNewline = go [] + where + go :: [WithLoc a] -> [WithLoc a] -> m [WithLoc a] + go stk tokens = case (stk, tokens) of + ([], []) -> return [] + (paren : _, []) -> throwLexicalErrorAt (loc paren) $ "unmatching paren found: " ++ show (value paren) + (_, token : tokens) | isOpen (value token) -> (token :) <$> go (token : stk) tokens + ([], token : _) | isClose (value token) -> throwLexicalErrorAt (loc token) $ "unmatching paren found: " ++ show (value token) + (_ : stk, token : tokens) | isClose (value token) -> (token :) <$> go stk tokens + (_ : _, token : tokens) | isNewline (value token) -> go stk tokens + (_, token : tokens) -> (token :) <$> go stk tokens + +removeEmptyLines :: forall a. (a -> Bool) -> [WithLoc a] -> [WithLoc a] +removeEmptyLines isNewline = go True + where + go :: Bool -> [WithLoc a] -> [WithLoc a] + go _ [] = [] + go lastIsNewline (token : tokens) + | lastIsNewline && isNewline (value token) = go True tokens + | otherwise = token : go (isNewline (value token)) tokens diff --git a/src/Jikka/Common/Parse/OffsideRule.hs b/src/Jikka/Common/Parse/OffsideRule.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/Parse/OffsideRule.hs @@ -0,0 +1,47 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Jikka.Common.Parse.OffsideRule + ( insertIndents, + ) +where + +import Jikka.Common.Error +import Jikka.Common.Location + +splitToLines :: forall a. (a -> Bool) -> [WithLoc a] -> [[WithLoc a]] +splitToLines isNewline = go [] + where + go :: [WithLoc a] -> [WithLoc a] -> [[WithLoc a]] + go [] [] = [] + go acc [] = [reverse acc] + go acc (token : tokens) + | isNewline (value token) = reverse (token : acc) : go [] tokens + | otherwise = go (token : acc) tokens + +insertIndents' :: forall m a. (MonadError Error m, Show a) => a -> a -> [[WithLoc a]] -> m [WithLoc a] +insertIndents' indent dedent = go [1] + where + go :: [Int] -> [[WithLoc a]] -> m [WithLoc a] + go stk tokens = case (stk, tokens) of + ([1], []) -> return [] + (_ : stk, []) -> (WithLoc (Loc 0 1 0) dedent :) <$> go stk [] + (_, [] : _) -> throwInternalError "a line must be non-empty" + (_, (token : _) : _) | column (loc token) < 0 -> throwInternalError $ "column must be 1-based for insertIndents': " ++ show token + ([], _) -> throwInternalError "too many dedents" + (x : stk, line@(token : _) : tokens') -> case compare x (column (loc token)) of + LT -> (withLoc (loc token) indent :) . (line ++) <$> go (column (loc token) : x : stk) tokens' + EQ -> (line ++) <$> go (x : stk) tokens' + GT -> case stk of + [] -> throwInternalError "too many dedents" + (x' : _) + | x' < column (loc token) -> throwLexicalErrorAt (loc token) $ "unindent does not match any outer indentation level: " ++ show token + | otherwise -> (withLoc (loc token) dedent :) <$> go stk (line : tokens') + withLoc :: Loc -> a -> WithLoc a + withLoc (Loc y x _) a = WithLoc (Loc y x 0) a + +-- | `insertIndents` inserts @INDENT@ and @DEDENT@ tokens with Python's way (). The `column` of `Loc` must be 1-based. This doen't use physical `line` of `Loc` because logical lines are used for indentation. +insertIndents :: forall m a. (MonadError Error m, Show a) => a -> a -> (a -> Bool) -> [WithLoc a] -> m [WithLoc a] +insertIndents indent dedent isNewline tokens = wrapError' "Jikka.Common.Parse.OffsideRule failed" $ do + let lines = splitToLines isNewline tokens + insertIndents' indent dedent lines diff --git a/src/Jikka/Common/Parse/Read.hs b/src/Jikka/Common/Parse/Read.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/Parse/Read.hs @@ -0,0 +1,7 @@ +module Jikka.Common.Parse.Read where + +import Data.Text (Text, unpack) +import Text.Read (readEither) + +run :: Read program => FilePath -> Text -> Either String program +run _ input = readEither $ unpack input diff --git a/src/Jikka/Common/Parse/ShuntingYard.hs b/src/Jikka/Common/Parse/ShuntingYard.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Common/Parse/ShuntingYard.hs @@ -0,0 +1,51 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Jikka.Common.Parse.ShuntingYard + ( run, + Prec, + Fixity (..), + BinOpInfo (..), + ) +where + +import Jikka.Common.Error +import Jikka.Common.Location + +type Prec = Int + +data Fixity + = Leftfix + | Rightfix + | Nonfix + deriving (Eq, Ord, Enum, Bounded, Show, Read) + +data BinOpInfo = BinOpInfo Fixity Prec + deriving (Eq, Ord, Show, Read) + +-- 10.6 Fixity Resolution - Haskell Language Report 2010 +-- https://www.haskell.org/onlinereport/haskell2010/haskellch10.html#x17-18100010.6 +run :: forall m op expr. MonadError Error m => (op -> m BinOpInfo) -> (WithLoc op -> WithLoc expr -> WithLoc expr -> WithLoc expr) -> (WithLoc expr, [(WithLoc op, WithLoc expr)]) -> m (WithLoc expr) +run info apply (e, tokens) = go [] [e] tokens + where + go :: [WithLoc op] -> [WithLoc expr] -> [(WithLoc op, WithLoc expr)] -> m (WithLoc expr) + go [] [e1] [] = return e1 + go (op : ops) (e2 : e1 : stk) [] = go ops (apply op e1 e2 : stk) [] + go [] stk ((op, e) : tokens) = go [op] (e : stk) tokens + go (op1 : ops) (e2 : e1 : stk) ((op2, e3) : tokens) = do + BinOpInfo fix1 prec1 <- info (value op1) + BinOpInfo fix2 prec2 <- info (value op2) + case () of + -- case (1): check for illegal expressions + _ + | prec1 == prec2 && (fix1 /= fix2 || fix1 == Nonfix) -> + throwSyntaxErrorAt (loc op1) "illigal expressions due to the fixity of operators" + -- case (2): op1 and op2 should associate to the left + _ + | prec1 > prec2 || (prec1 == prec2 && fix1 == Leftfix) -> + go ops (apply op1 e1 e2 : stk) ((op2, e3) : tokens) + -- case (3): op1 and op2 should associate to the right + _ + | otherwise -> + go (op2 : op1 : ops) (e3 : e2 : e1 : stk) tokens + go _ _ _ = throwInternalError "failed at shutting-yard algorithm" diff --git a/src/Jikka/Core/Convert.hs b/src/Jikka/Core/Convert.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert.hs @@ -0,0 +1,85 @@ +{-# LANGUAGE FlexibleContexts #-} + +-- | +-- Module : Jikka.Core.Convert +-- Description : is a module to combine other optimizers. / 他の最適化器を組み合わせて実行する module です。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- `Jikka.Core.Convert` is a module to combine other all optimizers. +module Jikka.Core.Convert + ( run, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import qualified Jikka.Core.Convert.Alpha as Alpha +import qualified Jikka.Core.Convert.ArithmeticalExpr as ArithmeticalExpr +import qualified Jikka.Core.Convert.Beta as Beta +import qualified Jikka.Core.Convert.BubbleLet as BubbleLet +import qualified Jikka.Core.Convert.CloseAll as CloseAll +import qualified Jikka.Core.Convert.CloseMin as CloseMin +import qualified Jikka.Core.Convert.CloseSum as CloseSum +import qualified Jikka.Core.Convert.ConstantFolding as ConstantFolding +import qualified Jikka.Core.Convert.ConstantPropagation as ConstantPropagation +import qualified Jikka.Core.Convert.ConvexHullTrick as ConvexHullTrick +import qualified Jikka.Core.Convert.CumulativeSum as CumulativeSum +import qualified Jikka.Core.Convert.Eta as Eta +import qualified Jikka.Core.Convert.MakeScanl as MakeScanl +import qualified Jikka.Core.Convert.MatrixExponentiation as MatrixExponentiation +import qualified Jikka.Core.Convert.PropagateMod as PropagateMod +import qualified Jikka.Core.Convert.RemoveUnusedVars as RemoveUnusedVars +import qualified Jikka.Core.Convert.SegmentTree as SegmentTree +import qualified Jikka.Core.Convert.ShortCutFusion as ShortCutFusion +import qualified Jikka.Core.Convert.SpecializeFoldl as SpecializeFoldl +import qualified Jikka.Core.Convert.StrengthReduction as StrengthReduction +import qualified Jikka.Core.Convert.TrivialLetElimination as TrivialLetElimination +import qualified Jikka.Core.Convert.TypeInfer as TypeInfer +import qualified Jikka.Core.Convert.UnpackTuple as UnpackTuple +import Jikka.Core.Language.Expr (Program) + +run'' :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run'' prog = do + prog <- RemoveUnusedVars.run prog + prog <- UnpackTuple.run prog + prog <- MatrixExponentiation.run prog + prog <- SpecializeFoldl.run prog + prog <- MakeScanl.run prog + prog <- PropagateMod.run prog + prog <- ConstantPropagation.run prog + prog <- ConstantFolding.run prog + prog <- ShortCutFusion.run prog + prog <- CloseSum.run prog + prog <- CloseAll.run prog + prog <- CloseMin.run prog + prog <- CumulativeSum.run prog + prog <- SegmentTree.run prog + prog <- BubbleLet.run prog + prog <- ArithmeticalExpr.run prog + prog <- ConvexHullTrick.run prog + prog <- StrengthReduction.run prog + Eta.run prog + +run' :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run' prog = do + prog <- Beta.run prog + prog <- TrivialLetElimination.run prog + prog <- run'' prog + prog <- run'' prog + prog <- run'' prog + prog <- run'' prog + run'' prog + +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = do + prog <- Alpha.run prog + prog <- TypeInfer.run prog + prog <- run' prog + prog <- run' prog + prog <- run' prog + prog <- run' prog + run' prog diff --git a/src/Jikka/Core/Convert/ANormal.hs b/src/Jikka/Core/Convert/ANormal.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/ANormal.hs @@ -0,0 +1,104 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.ANormal +-- Description : converts exprs to A-normal form. / 式を A 正規形に変換します。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Convert.ANormal + ( run, + ) +where + +import Jikka.Common.Alpha (MonadAlpha) +import Jikka.Common.Error +import qualified Jikka.Core.Convert.Alpha as Alpha (runProgram) +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Lint +import Jikka.Core.Language.TypeCheck +import Jikka.Core.Language.Util + +destruct :: (MonadAlpha m, MonadError Error m) => TypeEnv -> Expr -> m (TypeEnv, Expr -> Expr, Expr) +destruct env = \case + e@Var {} -> return (env, id, e) + e@Lit {} -> return (env, id, e) + e@App {} -> do + x <- genVarName' + t <- typecheckExpr env e + return ((x, t) : env, Let x t e, Var x) + e@Lam {} -> do + x <- genVarName' + t <- typecheckExpr env e + return ((x, t) : env, Let x t e, Var x) + Let x t e1 e2 -> do + (env, ctx, e1) <- destruct env e1 + (env, ctx', e2) <- destruct ((x, t) : env) e2 + return (env, ctx . Let x t e1 . ctx', e2) + +runApp :: (MonadAlpha m, MonadError Error m) => TypeEnv -> Expr -> [Expr] -> m Expr +runApp env f args = go env id args + where + go :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> ([Expr] -> [Expr]) -> [Expr] -> m Expr + go env acc [] = do + (_, ctx, f) <- destruct env f + return $ ctx (uncurryApp f (acc [])) + go env acc (arg : args) = do + (env, ctx, arg) <- destruct env arg + e <- go env (acc . (arg :)) args + return $ ctx e + +runExpr :: (MonadAlpha m, MonadError Error m) => TypeEnv -> Expr -> m Expr +runExpr env = \case + Var x -> return $ Var x + Lit lit -> return $ Lit lit + e@(App _ _) -> do + let (f, args) = curryApp e + f <- runExpr env f + args <- mapM (runExpr env) args + case (f, args) of + (Lit (LitBuiltin (If _)), [e1, e2, e3]) -> do + (_, ctx, e1) <- destruct env e1 + return $ ctx (App3 f e1 e2 e3) + _ -> runApp env f args + Lam x t body -> Lam x t <$> runExpr ((x, t) : env) body + Let x t e1 e2 -> do + e1 <- runExpr env e1 + (env, ctx, e1) <- destruct env e1 + e2 <- runExpr ((x, t) : env) e2 + return $ ctx (Let x t e1 e2) + +runToplevelExpr :: (MonadAlpha m, MonadError Error m) => TypeEnv -> ToplevelExpr -> m ToplevelExpr +runToplevelExpr env = \case + ResultExpr e -> ResultExpr <$> runExpr env e + ToplevelLet x t e cont -> do + e <- runExpr env e + cont <- runToplevelExpr ((x, t) : env) cont + return $ ToplevelLet x t e cont + ToplevelLetRec f args ret body cont -> do + let t = curryFunTy (map snd args) ret + body <- runExpr (reverse args ++ (f, t) : env) body + cont <- runToplevelExpr ((f, t) : env) cont + return $ ToplevelLetRec f args ret body cont + +-- | `run` makes a given program A-normal form. +-- A program is an A-normal form iff assigned exprs of all let-statements are values or function applications. +-- For example, this converts the following: +-- +-- > (let x = 1 in x) + ((fun y -> y) 1) +-- +-- to: +-- +-- > let x = 1 +-- > in let f = fun y -> y +-- > in let z = f x +-- > in z +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.ANormal" $ do + prog <- Alpha.runProgram prog + prog <- runToplevelExpr [] prog + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/Alpha.hs b/src/Jikka/Core/Convert/Alpha.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/Alpha.hs @@ -0,0 +1,83 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.Alpha +-- Description : does alpha-conversion. / alpha 変換をします。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Convert.Alpha where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Language.Expr + +rename :: MonadAlpha m => VarName -> m VarName +rename x = do + let base = takeWhile (/= '$') (unVarName x) + i <- nextCounter + return $ VarName (base ++ "$" ++ show i) + +runExpr :: (MonadAlpha m, MonadError Error m) => [(VarName, VarName)] -> Expr -> m Expr +runExpr env = \case + Var x -> case lookup x env of + Nothing -> throwInternalError $ "undefined variable: " ++ unVarName x + Just y -> return $ Var y + Lit lit -> return $ Lit lit + App f e -> App <$> runExpr env f <*> runExpr env e + Lam x t body -> do + y <- rename x + body <- runExpr ((x, y) : env) body + return $ Lam y t body + Let x t e1 e2 -> do + e1 <- runExpr env e1 + y <- rename x + e2 <- runExpr ((x, y) : env) e2 + return $ Let y t e1 e2 + +runToplevelExpr :: (MonadAlpha m, MonadError Error m) => [(VarName, VarName)] -> ToplevelExpr -> m ToplevelExpr +runToplevelExpr env = \case + ResultExpr e -> ResultExpr <$> runExpr env e + ToplevelLet x t e cont -> do + y <- rename x + e <- runExpr env e + cont <- runToplevelExpr ((x, y) : env) cont + return $ ToplevelLet y t e cont + ToplevelLetRec f args ret body cont -> do + g <- rename f + args <- forM args $ \(x, t) -> do + y <- rename x + return (x, y, t) + let args1 = map (\(x, y, _) -> (x, y)) args + let args2 = map (\(_, y, t) -> (y, t)) args + body <- runExpr (args1 ++ (f, g) : env) body + cont <- runToplevelExpr ((f, g) : env) cont + return $ ToplevelLetRec g args2 ret body cont + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram = runToplevelExpr [] + +-- | `run` renames variables in exprs to avoid name conflictions, even if the scopes of two variables are distinct. +-- +-- == Examples +-- +-- Before: +-- +-- > let x = 0 +-- > in y = x + x +-- > in x = x + y +-- > x + y +-- +-- After: +-- +-- > let x0 = 0 +-- > in y1 = x0 + x0 +-- > in x2 = x0 + y1 +-- > x2 + y1 +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.Alpha" $ do + runToplevelExpr [] prog diff --git a/src/Jikka/Core/Convert/ArithmeticalExpr.hs b/src/Jikka/Core/Convert/ArithmeticalExpr.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/ArithmeticalExpr.hs @@ -0,0 +1,51 @@ +{-# LANGUAGE FlexibleContexts #-} + +-- | +-- Module : Jikka.Core.Convert.ArithmeticalExpr +-- Description : sorts arithmetical exprs. / 算術式を整理します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Convert.ArithmeticalExpr + ( run, + ) +where + +import Jikka.Common.Error +import Jikka.Core.Language.ArithmeticalExpr +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Lint +import Jikka.Core.Language.TypeCheck +import Jikka.Core.Language.Util + +runExpr :: MonadError Error m => [(VarName, Type)] -> Expr -> m Expr +runExpr env e = do + t <- typecheckExpr env e + if t == IntTy + then return . formatArithmeticalExpr $ parseArithmeticalExpr e + else return e + +runProgram :: MonadError Error m => Program -> m Program +runProgram = mapExprProgramM runExpr -- Doesn't use RewriteRules because the rewriting may not terminate. + +-- | `run` sorts arithmetical exprs. +-- +-- == Examples +-- +-- Before: +-- +-- > 1 + a * 1 + b - b +-- +-- After: +-- +-- > a + 1 +run :: MonadError Error m => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.ArithmeticalExpr" $ do + precondition $ do + ensureWellTyped prog + prog <- runProgram prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/Beta.hs b/src/Jikka/Core/Convert/Beta.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/Beta.hs @@ -0,0 +1,55 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.Beta +-- Description : does beta-reductions. / beta 簡約を行います。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Convert.Beta + ( run, + + -- * internal rules + rule, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import qualified Jikka.Core.Convert.Alpha as Alpha +import Jikka.Core.Language.Beta +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Lint +import Jikka.Core.Language.RewriteRules + +rule :: MonadAlpha m => RewriteRule m +rule = RewriteRule $ \_ -> \case + App (Lam x _ e1) e2 -> Just <$> substitute x e2 e1 + _ -> return Nothing + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram = applyRewriteRuleProgram' rule + +-- | `run` does beta-reduction. +-- +-- == Examples +-- +-- Before: +-- +-- > (fun x -> x + x) y +-- +-- After: +-- +-- > y + y +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.Beta" $ do + precondition $ do + ensureWellTyped prog + prog <- Alpha.run prog + prog <- runProgram prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/BubbleLet.hs b/src/Jikka/Core/Convert/BubbleLet.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/BubbleLet.hs @@ -0,0 +1,62 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.BubbleLet +-- Description : makes let-exprs rise in higher-order functions. / 高階関数中の let 式を浮き上がらせます。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Convert.BubbleLet + ( run, + + -- * internal rules + rule, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.FreeVars +import Jikka.Core.Language.Lint +import Jikka.Core.Language.RewriteRules + +rule :: MonadAlpha m => RewriteRule m +rule = + let go f cont = case f of + Lam x t (Let y t' e body) | x `isUnusedVar` e -> return . Just $ Let y t' e (cont (Lam x t body)) + _ -> return Nothing + in RewriteRule $ \_ -> \case + Iterate' t k f x -> go f (\f -> Iterate' t k f x) + Foldl' t1 t2 f init xs -> go f (\f -> Foldl' t1 t2 f init xs) + Build' t f xs n -> go f (\f -> Build' t f xs n) + Map' t1 t2 f xs -> go f (\f -> Map' t1 t2 f xs) + Filter' t f xs -> go f (\f -> Filter' t f xs) + _ -> return Nothing + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram = applyRewriteRuleProgram' rule + +-- | `run` moves let-exprs in lambdas passed to higher-order functions to the outer of the higher-order functions. +-- +-- == Examples +-- +-- Before: +-- +-- > map (fun x -> let c = 12345 in c * x) xs +-- +-- After: +-- +-- > let c = 12345 in map (fun x -> c * x) xs +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.BubbleLet" $ do + precondition $ do + ensureWellTyped prog + prog <- runProgram prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/CloseAll.hs b/src/Jikka/Core/Convert/CloseAll.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/CloseAll.hs @@ -0,0 +1,142 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.CloseAll +-- Description : does reductions about @all@ and @any@, and tries to rewrite with closed-form exprs. / @all@ と @any@ についての簡約を行い、閉じた式への書き換えを目指します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Convert.CloseAll + ( run, + + -- * internal rules + rule, + reduceAll, + reduceAny, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Lint +import Jikka.Core.Language.RewriteRules +import Jikka.Core.Language.Util + +reduceAll :: MonadAlpha m => RewriteRule m +reduceAll = + let return' = return . Just + in RewriteRule $ \_ -> \case + -- list build functions + All' (Nil' _) -> return' LitTrue + All' (Cons' _ x xs) -> return' $ And' x (All' xs) + -- list map functions + All' (Reversed' _ xs) -> return' $ All' xs + All' (Sorted' _ xs) -> return' $ All' xs + All' (Filter' _ f xs) -> do + x <- genVarName' + return' $ All' (Map' BoolTy BoolTy (Lam x BoolTy (Implies' (App f (Var x)) (Var x))) xs) + All' (Map' _ _ f xs) -> case f of + Lam x _ (Not' e) -> do + return' $ Not' (Any' (Map' BoolTy BoolTy (Lam x BoolTy e) xs)) + Lam x _ (And' e1 e2) -> do + x1 <- genVarName x + x2 <- genVarName x + return' $ And' (All' (Map' BoolTy BoolTy (Lam x1 BoolTy e1) xs)) (All' (Map' BoolTy BoolTy (Lam x2 BoolTy e2) xs)) + _ -> return Nothing + -- others + _ -> return Nothing + +reduceAny :: MonadAlpha m => RewriteRule m +reduceAny = + let return' = return . Just + in RewriteRule $ \_ -> \case + -- list build functions + Any' (Nil' _) -> return' LitFalse + Any' (Cons' _ x xs) -> return' $ Or' x (Any' xs) + -- list map functions + Any' (Reversed' _ xs) -> return' $ Any' xs + Any' (Sorted' _ xs) -> return' $ Any' xs + Any' (Filter' _ f xs) -> do + x <- genVarName' + return' $ Any' (Map' BoolTy BoolTy (Lam x BoolTy (And' (App f (Var x)) (Var x))) xs) + Any' (Map' _ _ f xs) -> case f of + Lam x _ (Not' e) -> do + return' $ Not' (All' (Map' BoolTy BoolTy (Lam x BoolTy e) xs)) + Lam x _ (Or' e1 e2) -> do + x1 <- genVarName x + x2 <- genVarName x + return' $ Or' (Any' (Map' BoolTy BoolTy (Lam x1 BoolTy e1) xs)) (Any' (Map' BoolTy BoolTy (Lam x2 BoolTy e2) xs)) + Lam x _ (Implies' e1 e2) -> do + x1 <- genVarName x + x2 <- genVarName x + return' $ Or' (Any' (Map' BoolTy BoolTy (Lam x1 BoolTy (Negate' e1)) xs)) (Any' (Map' BoolTy BoolTy (Lam x2 BoolTy e2) xs)) + _ -> return Nothing + -- others + _ -> return Nothing + +rule :: MonadAlpha m => RewriteRule m +rule = + mconcat + [ reduceAll, + reduceAny + ] + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram = applyRewriteRuleProgram' rule + +-- | `run` reduces `All` and `Any`. +-- +-- == Examples +-- +-- Before: +-- +-- > any (filter (fun x -> x || f x) xs) +-- +-- After: +-- +-- > any xs || any (map f xs) +-- +-- == List of builtin functions which are reduced +-- +-- \[ +-- \newcommand\int{\mathbf{int}} +-- \newcommand\bool{\mathbf{bool}} +-- \newcommand\list{\mathbf{list}} +-- \] +-- +-- === Target functions +-- +-- * `All` \(: \list(\bool) \to \bool\) +-- * `Any` \(: \list(\bool) \to \bool\) +-- +-- === Boolean functions +-- +-- * `Not` \(: \bool \to \bool\) +-- * `And` \(: \bool \to \bool \to \bool\) +-- * `Or` \(: \bool \to \bool \to \bool\) +-- * `Implies` \(: \bool \to \bool \to \bool\) +-- +-- === List Build functions +-- +-- * `Nil` \(: \forall \alpha. \list(\alpha)\) +-- * `Cons` \(: \forall \alpha. \alpha \to \list(\alpha) \to \list(\alpha)\) +-- +-- === List Map functions +-- +-- * `Map` \(: \forall \alpha \beta. (\alpha \to \beta) \to \list(\alpha) \to \list(\beta)\) +-- * `Filter` \(: \forall \alpha \beta. (\alpha \to \bool) \to \list(\alpha) \to \list(\beta)\) +-- * `Reversed` \(: \forall \alpha. \list(\alpha) \to \list(\alpha)\) +-- * `Sorted` \(: \forall \alpha. \list(\alpha) \to \list(\alpha)\) +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.CloseAll" $ do + precondition $ do + ensureWellTyped prog + prog <- runProgram prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/CloseMin.hs b/src/Jikka/Core/Convert/CloseMin.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/CloseMin.hs @@ -0,0 +1,171 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.CloseMin +-- Description : does reductions about minnimums and maximums of lists, and tries to rewrite with closed-form exprs. / リストの最小値と最大値についての簡約を行い、閉じた式への書き換えを目指します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Convert.CloseMin + ( run, + + -- * internal rules + rule, + reduceMin, + reduceMax, + reduceArgMin, + reduceArgMax, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import qualified Jikka.Core.Convert.Alpha as Alpha +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.FreeVars +import Jikka.Core.Language.Lint +import Jikka.Core.Language.RewriteRules + +reduceMin :: Monad m => RewriteRule m +reduceMin = simpleRewriteRule $ \case + -- list build functions + Min1' t (Nil' _) -> Just $ Bottom' t "no minimum in empty list" + Min1' _ (Cons' _ e (Nil' _)) -> Just e + Min1' t (Cons' _ e (Cons' _ e' es)) -> Just $ Min2' t e (Min1' t (Cons' t e' es)) + -- list map functions + Min1' t (Reversed' _ es) -> Just $ Min1' t es + Min1' t (Cons' _ e (Reversed' _ es)) -> Just $ Min1' t (Cons' t e es) + Min1' t (Sorted' _ es) -> Just $ Min1' t es + Min1' t (Cons' _ e (Sorted' _ es)) -> Just $ Min1' t (Cons' t e es) + Min1' t (Map' t1 t2 f es) -> case f of + Lam x _ e | x `isUnusedVar` e -> Just e + Lam x _ (Min2' _ e1 e2) -> Just $ Min2' t (Min1' t (Map' t1 t2 (Lam x t e1) es)) (Min1' t (Map' t1 t2 (Lam x t e2) es)) + Lam x _ (Negate' e) -> Just $ Negate' (Max1' t (Map' t1 t2 (Lam x IntTy e) es)) + Lam x _ (Plus' e1 e2) | x `isUnusedVar` e1 -> Just $ Plus' e1 (Min1' t (Map' t1 t2 (Lam x IntTy e2) es)) + Lam x _ (Plus' e1 e2) | x `isUnusedVar` e2 -> Just $ Plus' (Min1' t (Map' t1 t2 (Lam x IntTy e1) es)) e2 + _ -> Nothing + Min1' t (Cons' _ e0 (Map' t1 t2 f xs)) -> case f of + Lam x _ e | x `isUnusedVar` e -> Just $ If' t (Equal' IntTy (Len' t xs) Lit0) e0 (Min2' t e0 e) + Lam x _ (Min2' _ e1 e2) -> Just $ Min2' t (Min1' t (Cons' t e0 (Map' t1 t2 (Lam x t e1) xs))) (Min1' t (Cons' t e0 (Map' t1 t2 (Lam x t e2) xs))) + Lam x _ (Negate' e) -> Just $ Negate' (Max1' t (Cons' t (Negate' e0) (Map' t1 t2 (Lam x IntTy e) xs))) + Lam x _ (Plus' e1 e2) | x `isUnusedVar` e1 -> Just $ Plus' e1 (Min1' t (Cons' t (Minus' e0 e1) (Map' t1 t2 (Lam x IntTy e2) xs))) + Lam x _ (Plus' e1 e2) | x `isUnusedVar` e2 -> Just $ Plus' (Min1' t (Cons' t (Minus' e0 e1) (Map' t1 t2 (Lam x IntTy e1) xs))) e2 + _ -> Nothing + _ -> Nothing + +reduceMax :: Monad m => RewriteRule m +reduceMax = simpleRewriteRule $ \case + -- list build functions + Max1' t (Nil' _) -> Just $ Bottom' t "no maximum in empty list" + Max1' _ (Cons' _ e (Nil' _)) -> Just e + Max1' t (Cons' _ e (Cons' _ e' es)) -> Just $ Max2' t e (Max1' t (Cons' t e' es)) + -- list map functions + Max1' t (Reversed' _ es) -> Just $ Max1' t es + Max1' t (Cons' _ e (Reversed' _ es)) -> Just $ Max1' t (Cons' t e es) + Max1' t (Sorted' _ es) -> Just $ Max1' t es + Max1' t (Cons' _ e (Sorted' _ es)) -> Just $ Max1' t (Cons' t e es) + Max1' t (Map' t1 t2 f es) -> case f of + Lam x _ e | x `isUnusedVar` e -> Just e + Lam x _ (Max2' _ e1 e2) -> Just $ Max2' t (Map' t1 t2 (Lam x t e1) es) (Map' t1 t2 (Lam x t e2) es) + Lam x _ (Negate' e) -> Just $ Negate' (Min1' t2 (Map' t1 t2 (Lam x IntTy e) es)) + Lam x _ (Plus' e1 e2) | x `isUnusedVar` e1 -> Just $ Plus' e1 (Max1' t2 (Map' t1 t2 (Lam x IntTy e2) es)) + Lam x _ (Plus' e1 e2) | x `isUnusedVar` e2 -> Just $ Plus' (Max1' t2 (Map' t1 t2 (Lam x IntTy e1) es)) e2 + _ -> Nothing + Max1' t (Cons' _ e0 (Map' t1 t2 f xs)) -> case f of + Lam x _ e | x `isUnusedVar` e -> Just $ If' t (Equal' IntTy (Len' t xs) Lit0) e0 (Max2' t e0 e) + Lam x _ (Max2' _ e1 e2) -> Just $ Max2' t (Max1' t (Cons' t e0 (Map' t1 t2 (Lam x t e1) xs))) (Max1' t (Cons' t e0 (Map' t1 t2 (Lam x t e2) xs))) + Lam x _ (Negate' e) -> Just $ Negate' (Min1' t (Cons' t (Negate' e0) (Map' t1 t2 (Lam x IntTy e) xs))) + Lam x _ (Plus' e1 e2) | x `isUnusedVar` e1 -> Just $ Plus' e1 (Max1' t (Cons' t (Minus' e0 e1) (Map' t1 t2 (Lam x IntTy e2) xs))) + Lam x _ (Plus' e1 e2) | x `isUnusedVar` e2 -> Just $ Plus' (Max1' t (Cons' t (Minus' e0 e1) (Map' t1 t2 (Lam x IntTy e1) xs))) e2 + _ -> Nothing + _ -> Nothing + +-- | TODO: implement this +reduceArgMin :: Monad m => RewriteRule m +reduceArgMin = simpleRewriteRule $ \case + -- list map functions + ArgMin' t (Reversed' _ xs) -> Just $ Minus' (Minus' (Len' t xs) (ArgMin' t xs)) Lit1 + ArgMin' _ (Map' _ _ (Lam x _ e) _) | x `isUnusedVar` e -> Just Lit0 + ArgMin' _ (Map' t1 t2 (Lam x t (Plus' e1 e2)) xs) | x `isUnusedVar` e1 -> Just $ ArgMin' t2 (Map' t1 t2 (Lam x t e2) xs) + ArgMin' _ (Map' t1 t2 (Lam x t (Plus' e1 e2)) xs) | x `isUnusedVar` e2 -> Just $ ArgMin' t2 (Map' t1 t2 (Lam x t e1) xs) + _ -> Nothing + +-- | TODO: implement this +reduceArgMax :: Monad m => RewriteRule m +reduceArgMax = simpleRewriteRule $ \case + -- list map functions + ArgMax' t (Reversed' _ xs) -> Just $ Minus' (Minus' (Len' t xs) (ArgMax' t xs)) Lit1 + ArgMax' _ (Map' _ _ (Lam x t e) xs) | x `isUnusedVar` e -> Just $ Minus' (Len' t xs) Lit1 + ArgMax' _ (Map' t1 t2 (Lam x t (Plus' e1 e2)) xs) | x `isUnusedVar` e1 -> Just $ ArgMax' t2 (Map' t1 t2 (Lam x t e2) xs) + ArgMax' _ (Map' t1 t2 (Lam x t (Plus' e1 e2)) xs) | x `isUnusedVar` e2 -> Just $ ArgMax' t2 (Map' t1 t2 (Lam x t e1) xs) + _ -> Nothing + +rule :: Monad m => RewriteRule m +rule = + mconcat + [ reduceMin, + reduceMax, + reduceArgMin, + reduceArgMax + ] + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram = applyRewriteRuleProgram' rule + +-- | `run` reduces maximums and minimums. +-- +-- == Examples +-- +-- Before: +-- +-- > max (map (fun x -> 3 + f x) xs) +-- +-- After: +-- +-- > 3 + max (map f xs) +-- +-- == List of builtin functions which are reduced +-- +-- \[ +-- \newcommand\int{\mathbf{int}} +-- \newcommand\bool{\mathbf{bool}} +-- \newcommand\list{\mathbf{list}} +-- \] +-- +-- === Target functions +-- +-- * `Max1` \(: \forall \alpha. \list(\alpha) \to \alpha\) +-- * `Min1` \(: \forall \alpha. \list(\alpha) \to \alpha\) +-- * `ArgMax` \(: \forall \alpha. \list(\alpha) \to \int\) +-- * `ArgMin` \(: \forall \alpha. \list(\alpha) \to \int\) +-- +-- === Related functions +-- +-- * `Max2` \(: \forall \alpha. \alpha \to \alpha \to \alpha\) +-- * `Min2` \(: \forall \alpha. \alpha \to \alpha \to \alpha\) +-- +-- === List Build functions +-- +-- * `Nil` \(: \forall \alpha. \list(\alpha)\) +-- * `Cons` \(: \forall \alpha. \alpha \to \list(\alpha) \to \list(\alpha)\) +-- * `Range1` \(: \int \to \list(\int)\) +-- +-- === List Map functions +-- +-- * `Map` \(: \forall \alpha \beta. (\alpha \to \beta) \to \list(\alpha) \to \list(\beta)\) +-- * `Filter` \(: \forall \alpha \beta. (\alpha \to \bool) \to \list(\alpha) \to \list(\beta)\) +-- * `Reversed` \(: \forall \alpha. \list(\alpha) \to \list(\alpha)\) +-- * `Sorted` \(: \forall \alpha. \list(\alpha) \to \list(\alpha)\) +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.CloseMin" $ do + precondition $ do + ensureWellTyped prog + prog <- runProgram prog + prog <- Alpha.run prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/CloseSum.hs b/src/Jikka/Core/Convert/CloseSum.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/CloseSum.hs @@ -0,0 +1,212 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.CloseSum +-- Description : does reductions about summations and products, and tries to rewrite with closed-form exprs. / 総和と総乗についての簡約を行い、閉じた式への書き換えを目指します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- \[ +-- \newcommand\int{\mathbf{int}} +-- \newcommand\bool{\mathbf{bool}} +-- \newcommand\list{\mathbf{list}} +-- \] +module Jikka.Core.Convert.CloseSum + ( run, + + -- * internal rules + rule, + reduceSum, + reduceProduct, + reduceModSum, + reduceModProduct, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.FreeVars +import Jikka.Core.Language.Lint +import Jikka.Core.Language.RewriteRules +import Jikka.Core.Language.Util + +reduceSum :: MonadAlpha m => RewriteRule m +reduceSum = + let return' = return . Just + in RewriteRule $ \_ -> \case + Sum' xs -> case xs of + -- reduce list build functions + Nil' _ -> return' Lit0 + Cons' _ x xs -> return' $ Plus' x (Sum' xs) + Range1' n -> return' $ FloorDiv' (Mult' n (Minus' n Lit1)) Lit2 + -- reduce list map functions + Reversed' _ xs -> return' $ Sum' xs + Sorted' _ xs -> return' $ Sum' xs + Filter' _ g (Map' t1 _ f xs) -> do + x <- genVarName' + let h = Lam x t1 (If' IntTy (App g (App f (Var x))) (App f (Var x)) Lit0) + return' $ Sum' (Map' t1 IntTy h xs) + Map' t1 IntTy (Lam x _ body) xs -> case (body, xs) of + (e, xs) | x `isUnusedVar` e -> return' $ Mult' (Len' t1 xs) e + (body, Range1' n) | body == Var x -> return' $ FloorDiv' (Mult' n (Minus' n Lit1)) Lit2 + (body, Range1' n) | body == Mult' (Var x) (Var x) || body == Pow' (Var x) (LitInt' 2) -> return' $ FloorDiv' (Mult' n (Mult' (Minus' n Lit1) (Minus' (Mult' Lit2 (Var x)) Lit1))) (LitInt' 6) + (Negate' e, xs) -> return' $ Negate' (Sum' (Map' t1 IntTy (Lam x t1 e) xs)) + (Plus' e1 e2, xs) -> return' $ Plus' (Sum' (Map' t1 IntTy (Lam x t1 e1) xs)) (Sum' (Map' t1 IntTy (Lam x t1 e2) xs)) + (Minus' e1 e2, xs) -> return' $ Minus' (Sum' (Map' t1 IntTy (Lam x t1 e1) xs)) (Sum' (Map' t1 IntTy (Lam x t1 e2) xs)) + (Mult' e1 e2, xs) | x `isUnusedVar` e1 -> return' $ Mult' e1 (Sum' (Map' t1 IntTy (Lam x t1 e2) xs)) + (Mult' e1 e2, xs) | x `isUnusedVar` e2 -> return' $ Mult' e2 (Sum' (Map' t1 IntTy (Lam x t1 e1) xs)) + _ -> return Nothing + -- others + _ -> return Nothing + _ -> return Nothing + +-- | TODO: implement this. +reduceProduct :: Monad m => RewriteRule m +reduceProduct = simpleRewriteRule $ \case + Product' xs -> case xs of + -- reduce list build functions + Nil' _ -> Just Lit1 + Cons' _ x xs -> Just $ Mult' x (Product' xs) + Range1' n -> Just $ If' IntTy (Equal' IntTy n Lit0) Lit1 Lit0 + -- reduce list map functions + Reversed' _ xs -> Just $ Product' xs + Sorted' _ xs -> Just $ Product' xs + Map' t1 _ (Lam x _ e) xs | x `isUnusedVar` e -> Just $ Pow' e (Len' t1 xs) + Map' t1 t2 (Lam x t (Negate' e)) xs -> Just $ Mult' (Pow' (Negate' Lit0) (Len' t1 xs)) (Product' (Map' t1 t2 (Lam x t e) xs)) + Map' t1 t2 (Lam x t (Mult' e1 e2)) xs -> Just $ Mult' (Product' (Map' t1 t2 (Lam x t e1) xs)) (Product' (Map' t1 t2 (Lam x t e2) xs)) + -- others + _ -> Nothing + _ -> Nothing + +-- | +-- * This assumes that `ModFloor` is already propagated. +reduceModSum :: MonadAlpha m => RewriteRule m +reduceModSum = + let return' = return . Just + in RewriteRule $ \_ -> \case + ModSum' xs m -> case xs of + -- the corner case + _ | m == Lit1 -> return' Lit0 + -- reduce list build functions + Nil' _ -> return' Lit0 + Cons' _ x xs -> return' $ ModPlus' x (ModSum' xs m) m + Range1' n -> return' $ ModMult' (ModMult' n (ModPlus' n Lit1 m) m) (ModInv' Lit2 m) m + -- reduce list map functions + Reversed' _ xs -> return' $ ModSum' xs m + Sorted' _ xs -> return' $ ModSum' xs m + Filter' _ g (Map' t1 _ f xs) -> do + x <- genVarName' + let h = Lam x t1 (If' IntTy (App g (App f (Var x))) (App f (Var x)) Lit0) + return' $ ModSum' (Map' t1 IntTy h xs) m + Map' t1 IntTy (Lam x _ body) xs -> do + let go body = case (body, xs) of + (e, xs) | x `isUnusedVar` e -> return' $ ModMult' (Len' t1 xs) e m + (body, Range1' n) | body == Var x -> return' $ ModMult' (ModMult' n (ModMinus' n Lit1 m) m) (ModInv' Lit2 m) m + (body, Range1' n) | body == ModMult' (Var x) (Var x) m || body == ModPow' (Var x) (LitInt' 2) m -> return' $ ModMult' (ModMult' n (ModMult' (ModMinus' n Lit1 m) (ModMinus' (ModMult' Lit2 n m) Lit1 m) m) m) (ModInv' (LitInt' 6) m) m + (ModNegate' e m', xs) | m' == m -> return' $ ModNegate' (ModSum' (Map' t1 IntTy (Lam x t1 e) xs) m) m + (ModPlus' e1 e2 m', xs) | m' == m -> return' $ ModPlus' (ModSum' (Map' t1 IntTy (Lam x t1 e1) xs) m) (ModSum' (Map' t1 IntTy (Lam x t1 e2) xs) m) m + (ModMinus' e1 e2 m', xs) | m' == m -> return' $ ModMinus' (ModSum' (Map' t1 IntTy (Lam x t1 e1) xs) m) (ModSum' (Map' t1 IntTy (Lam x t1 e2) xs) m) m + (ModMult' e1 e2 m', xs) | x `isUnusedVar` e1 && m' == m -> return' $ ModMult' e1 (ModSum' (Map' t1 IntTy (Lam x t1 e2) xs) m) m + (ModMult' e1 e2 m', xs) | x `isUnusedVar` e2 && m' == m -> return' $ ModMult' e2 (ModSum' (Map' t1 IntTy (Lam x t1 e1) xs) m) m + _ -> return Nothing + case body of + FloorMod' body m' | m' == m -> go body -- We shouldn't remove FloorMod not to introduce loops of rewrite rules between Jikka.Core.Convert.PropagateMod. + _ -> go body + -- others + _ -> return Nothing + _ -> return Nothing + +-- | TODO: implement this. +reduceModProduct :: Monad m => RewriteRule m +reduceModProduct = simpleRewriteRule $ \case + ModProduct' xs m -> case xs of + -- the corner case + _ | m == Lit1 -> Just Lit0 + -- reduce list build functions + Nil' _ -> Just $ FloorMod' Lit1 m + Cons' _ x xs -> Just $ ModMult' x (ModProduct' xs m) m + Range1' n -> Just $ If' IntTy (Equal' IntTy n Lit0) (FloorMod' Lit1 m) Lit0 + -- reduce list map functions + Reversed' _ xs -> Just $ ModProduct' xs m + Sorted' _ xs -> Just $ ModProduct' xs m + -- others + _ -> Nothing + _ -> Nothing + +rule :: MonadAlpha m => RewriteRule m +rule = + mconcat + [ reduceSum, + reduceProduct, + reduceModSum, + reduceModProduct + ] + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram = applyRewriteRuleProgram' rule + +-- | `run` reduces summations and products. +-- +-- * This doen't do nothing about `Foldl`. +-- +-- == Examples +-- +-- Before: +-- +-- > foldl (fun y x -> y + x) 0 (range n) +-- +-- After: +-- +-- > x * (x - 1) / 2 +-- +-- == List of builtin functions which are reduced +-- +-- === Target functions +-- +-- * `Sum` \(: \list(\int) \to \int\) +-- * `Product` \(: \list(\int) \to \int\) +-- * `ModSum` \(: \list(\int) \to \int \to \int\) +-- * `ModProduct` \(: \list(\int) \to \int \to \int\) +-- +-- === Arithmetical functions +-- +-- * `Negate` \(: \int \to \int\) +-- * `Plus` \(: \int \to \int \to \int\) +-- * `Minus` \(: \int \to \int \to \int\) +-- * `Mult` \(: \int \to \int \to \int\) +-- * `Pow` \(: \int \to \int \to \int\) +-- +-- === Arithmetical functions with modulo +-- +-- * `ModNegate` \(: \int \to \int \to \int\) +-- * `ModPlus` \(: \int \to \int \to \int \to \int\) +-- * `ModMinus` \(: \int \to \int \to \int \to \int\) +-- * `ModMult` \(: \int \to \int \to \int \to \int\) +-- * `ModPow` \(: \int \to \int \to \int \to \int\) +-- +-- === List Build functions +-- +-- * `Nil` \(: \forall \alpha. \list(\alpha)\) +-- * `Cons` \(: \forall \alpha. \alpha \to \list(\alpha) \to \list(\alpha)\) +-- * `Range1` \(: \int \to \list(\int)\) +-- +-- === List Map functions +-- +-- * `Map` \(: \forall \alpha \beta. (\alpha \to \beta) \to \list(\alpha) \to \list(\beta)\) +-- * `Filter` \(: \forall \alpha \beta. (\alpha \to \bool) \to \list(\alpha) \to \list(\beta)\) +-- * `Reversed` \(: \forall \alpha. \list(\alpha) \to \list(\alpha)\) +-- * `Sorted` \(: \forall \alpha. \list(\alpha) \to \list(\alpha)\) +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.CloseSum" $ do + precondition $ do + ensureWellTyped prog + prog <- runProgram prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/ConstantFolding.hs b/src/Jikka/Core/Convert/ConstantFolding.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/ConstantFolding.hs @@ -0,0 +1,237 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.ConstantFolding +-- Description : folds constants. / 定数畳み込みをします。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- \[ +-- \newcommand\int{\mathbf{int}} +-- \newcommand\bool{\mathbf{bool}} +-- \newcommand\list{\mathbf{list}} +-- \] +module Jikka.Core.Convert.ConstantFolding + ( run, + + -- * internal rules + rule, + reduceConstArithmeticalExpr, + reduceConstMaxExpr, + reduceConstBooleanExpr, + reduceConstBitExpr, + reduceConstComparison, + ) +where + +import Data.Bits +import Data.Either +import Jikka.Common.Error +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Lint +import Jikka.Core.Language.RewriteRules +import Jikka.Core.Language.Runtime + +-- | +-- == List of functions which are reduced +-- +-- === Basic arithmetical functions +-- +-- * `Negate` \(: \int \to \int\) +-- * `Plus` \(: \int \to \int \to \int\) +-- * `Minus` \(: \int \to \int \to \int\) +-- * `Mult` \(: \int \to \int \to \int\) +-- * `FloorDiv` \(: \int \to \int \to \int\) +-- * `FloorMod` \(: \int \to \int \to \int\) +-- * `CeilDiv` \(: \int \to \int \to \int\) +-- * `CeilMod` \(: \int \to \int \to \int\) +-- * `Pow` \(: \int \to \int \to \int\) +-- +-- === Advanced arithmetical functions +-- +-- * `Abs` \(: \int \to \int\) +-- * `Gcd` \(: \int \to \int \to \int\) +-- * `Lcm` \(: \int \to \int \to \int\) +reduceConstArithmeticalExpr :: Monad m => RewriteRule m +reduceConstArithmeticalExpr = + let return' = Just . LitInt' + in simpleRewriteRule $ \case + Negate' (LitInt' a) -> return' $ - a + Plus' a (LitInt' 0) -> Just a + Plus' (LitInt' 0) b -> Just b + Plus' (LitInt' a) (LitInt' b) -> return' $ a + b + Minus' a (LitInt' 0) -> Just a + Minus' (LitInt' 0) b -> Just (Negate' b) + Minus' (LitInt' a) (LitInt' b) -> return' $ a - b + Mult' _ (LitInt' 0) -> return' 0 + Mult' a (LitInt' 1) -> Just a + Mult' (LitInt' 0) _ -> return' 0 + Mult' (LitInt' 1) b -> Just b + Mult' (LitInt' a) (LitInt' b) -> return' $ a * b + FloorDiv' a (LitInt' 1) -> Just a + FloorDiv' (LitInt' a) (LitInt' b) -> Just . fromRight (Bottom' IntTy "division by zero") . (LitInt' <$>) $ floorDiv a b + FloorMod' _ (LitInt' 1) -> return' 0 + FloorMod' (LitInt' a) (LitInt' b) -> Just . fromRight (Bottom' IntTy "modulo by zero") . (LitInt' <$>) $ floorMod a b + CeilDiv' a (LitInt' 1) -> Just a + CeilDiv' (LitInt' a) (LitInt' b) -> Just . fromRight (Bottom' IntTy "division by zero") . (LitInt' <$>) $ ceilDiv a b + CeilMod' _ (LitInt' 1) -> return' 0 + CeilMod' (LitInt' a) (LitInt' b) -> Just . fromRight (Bottom' IntTy "modulo by zero") . (LitInt' <$>) $ ceilMod a b + Pow' _ (LitInt' 0) -> return' 1 + Pow' a (LitInt' 1) -> Just a + Pow' (LitInt' a) (LitInt' b) | b >= 0 && fromInteger b * log (abs (fromInteger a)) < 100 -> return' $ a ^ b + Abs' (LitInt' a) -> return' $ abs a + Gcd' a (LitInt' 0) -> Just a + Gcd' _ (LitInt' 1) -> return' 1 + Gcd' (LitInt' 0) b -> Just b + Gcd' (LitInt' 1) _ -> return' 1 + Gcd' (LitInt' a) (LitInt' b) -> return' $ gcd a b + Lcm' _ (LitInt' 0) -> return' 0 + Lcm' a (LitInt' 1) -> Just a + Lcm' (LitInt' 0) _ -> return' 0 + Lcm' (LitInt' 1) b -> Just b + Lcm' (LitInt' a) (LitInt' b) -> return' $ lcm a b + _ -> Nothing + +-- | +-- == List of functions which are reduced +-- +-- === Max functions +-- +-- * `Min2` \(: \forall \alpha. \alpha \to \alpha \to \alpha\) (specialized to \(\alpha = \lbrace \bool, \int \rbrace\)) +-- * `Max2` \(: \forall \alpha. \alpha \to \alpha \to \alpha\) (specialized to \(\alpha = \lbrace \bool, \int \rbrace\)) +reduceConstMaxExpr :: Monad m => RewriteRule m +reduceConstMaxExpr = simpleRewriteRule $ \case + Min2' _ (LitInt' a) (LitInt' b) -> Just . LitInt' $ min a b + Min2' _ (LitBool' a) (LitBool' b) -> Just . LitBool' $ min a b + Max2' _ (LitInt' a) (LitInt' b) -> Just . LitInt' $ max a b + Max2' _ (LitBool' a) (LitBool' b) -> Just . LitBool' $ max a b + _ -> Nothing + +-- | +-- == List of functions which are reduced +-- +-- === Boolean functions +-- +-- * `Not` \(: \bool \to \bool\) +-- * `And` \(: \bool \to \bool \to \bool\) +-- * `Or` \(: \bool \to \bool \to \bool\) +-- * `Implies` \(: \bool \to \bool \to \bool\) +-- * `If` \(: \forall \alpha. \bool \to \alpha \to \alpha \to \alpha\) +reduceConstBooleanExpr :: Monad m => RewriteRule m +reduceConstBooleanExpr = simpleRewriteRule $ \case + Not' (LitBool' a) -> Just $ LitBool' (not a) + And' _ LitFalse -> Just LitFalse + And' a LitTrue -> Just a + And' LitFalse _ -> Just LitFalse + And' LitTrue b -> Just b + Or' a LitFalse -> Just a + Or' _ LitTrue -> Just LitTrue + Or' LitFalse b -> Just b + Or' LitTrue _ -> Just LitTrue + Implies' a LitFalse -> Just $ Not' a + Implies' _ LitTrue -> Just LitTrue + Implies' LitFalse _ -> Just LitTrue + Implies' LitTrue a -> Just a + If' _ (LitBool' a) e1 e2 -> Just $ if a then e1 else e2 + _ -> Nothing + +-- | +-- == List of functions which are reduced +-- +-- === Bitwise boolean functions +-- +-- * `BitNot` \(: \int \to \int\) +-- * `BitAnd` \(: \int \to \int \to \int\) +-- * `BitOr` \(: \int \to \int \to \int\) +-- * `BitXor` \(: \int \to \int \to \int\) +-- * `BitLeftShift` \(: \int \to \int \to \int\) +-- * `BitRightShift` \(: \int \to \int \to \int\) +reduceConstBitExpr :: Monad m => RewriteRule m +reduceConstBitExpr = + let return' = Just . LitInt' + in simpleRewriteRule $ \case + BitNot' (LitInt' a) -> return' $ complement a + BitAnd' _ (LitInt' 0) -> return' 0 + BitAnd' a (LitInt' (-1)) -> Just a + BitAnd' (LitInt' 0) _ -> return' 0 + BitAnd' (LitInt' (-1)) b -> Just b + BitAnd' (LitInt' a) (LitInt' b) -> return' $ a .&. b + BitOr' a (LitInt' 0) -> Just a + BitOr' _ (LitInt' (-1)) -> return' $ -1 + BitOr' (LitInt' 0) b -> Just b + BitOr' (LitInt' (-1)) _ -> return' $ -1 + BitOr' (LitInt' a) (LitInt' b) -> return' $ a .|. b + BitXor' a (LitInt' 0) -> Just a + BitXor' a (LitInt' (-1)) -> Just $ BitNot' a + BitXor' (LitInt' 0) b -> Just b + BitXor' (LitInt' (-1)) b -> Just $ BitNot' b + BitXor' (LitInt' a) (LitInt' b) -> return' $ a `xor` b + BitLeftShift' a (LitInt' 0) -> Just a + BitLeftShift' (LitInt' 0) _ -> return' 0 + BitLeftShift' (LitInt' a) (LitInt' b) | - 100 < b && b < 100 -> return' $ a `shift` fromInteger b + BitRightShift' a (LitInt' 0) -> Just a + BitRightShift' (LitInt' 0) _ -> return' 0 + BitRightShift' (LitInt' a) (LitInt' b) | - 100 < b && b < 100 -> return' $ a `shift` fromInteger (- b) + _ -> Nothing + +-- | +-- == List of functions which are reduced +-- +-- === Comparison functions +-- +-- * `LessThan` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\)) +-- * `LessEqual` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\)) +-- * `GreaterThan` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\)) +-- * `GreaterEqual` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\)) +-- * `Equal` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\)) +-- * `NotEqual` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\)) +reduceConstComparison :: Monad m => RewriteRule m +reduceConstComparison = + simpleRewriteRule $ + (LitBool' <$>) . \case + LessThan' _ (LitInt' a) (LitInt' b) -> Just $ a < b + LessEqual' _ (LitBool' a) (LitBool' b) -> Just $ a <= b + LessEqual' _ (LitInt' a) (LitInt' b) -> Just $ a <= b + GreaterThan' _ (LitBool' a) (LitBool' b) -> Just $ a > b + GreaterThan' _ (LitInt' a) (LitInt' b) -> Just $ a > b + GreaterEqual' _ (LitBool' a) (LitBool' b) -> Just $ a >= b + Equal' _ (LitInt' a) (LitInt' b) -> Just $ a == b + Equal' _ (LitBool' a) (LitBool' b) -> Just $ a == b + NotEqual' _ (LitInt' a) (LitInt' b) -> Just $ a /= b + NotEqual' _ (LitBool' a) (LitBool' b) -> Just $ a /= b + _ -> Nothing + +rule :: MonadError Error m => RewriteRule m +rule = + mconcat + [ reduceConstArithmeticalExpr, + reduceConstMaxExpr, + reduceConstBooleanExpr, + reduceConstBitExpr, + reduceConstComparison + ] + +runProgram :: MonadError Error m => Program -> m Program +runProgram = applyRewriteRuleProgram' rule + +-- | `run` folds constants in given programs. +-- For example, this converts the following: +-- +-- > 3 x + 2 + 1 +-- +-- to the follwoing: +-- +-- > 3 x + 3 +run :: MonadError Error m => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.ConstantFolding" $ do + precondition $ do + ensureWellTyped prog + prog <- runProgram prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/ConstantPropagation.hs b/src/Jikka/Core/Convert/ConstantPropagation.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/ConstantPropagation.hs @@ -0,0 +1,75 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.ConstantPropagation +-- Description : propagates something constants, for exprs which are computable with constant time. / 定数時間で計算できるような式についての、ある種の定数伝播をします。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Convert.ConstantPropagation + ( run, + run', + ) +where + +import qualified Data.Map as M +import Data.Maybe (fromMaybe) +import Jikka.Common.Error +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Lint +import Jikka.Core.Language.Util + +type Env = M.Map VarName Expr + +runExpr :: Env -> Expr -> Expr +runExpr env = \case + Var x -> fromMaybe (Var x) (M.lookup x env) + Lit lit -> Lit lit + App f e -> App (runExpr env f) (runExpr env e) + Lam x t body -> Lam x t (runExpr env body) + Let x t e1 e2 -> + let e1' = runExpr env e1 + in if isConstantTimeExpr e1' + then runExpr (M.insert x e1' env) e2 + else Let x t e1' (runExpr env e2) + +runToplevelExpr :: Env -> ToplevelExpr -> ToplevelExpr +runToplevelExpr env = \case + ResultExpr e -> ResultExpr (runExpr env e) + ToplevelLet x t e cont -> + let e' = runExpr env e + in if isConstantTimeExpr e' + then runToplevelExpr (M.insert x e' env) cont + else ToplevelLet x t e' (runToplevelExpr env cont) + ToplevelLetRec f args ret body cont -> + ToplevelLetRec f args ret (runExpr env body) (runToplevelExpr env cont) + +run' :: Program -> Program +run' = runToplevelExpr M.empty + +-- | `run` does constant propagation. +-- This assumes that the program is alpha-converted. +-- +-- For example, this converts the following: +-- +-- > let x = 1 +-- > in let f = fun y -> y +-- > in x + x + f(x) +-- +-- to: +-- +-- > let f = fun y -> y +-- > in 1 + 1 + f(1) +-- +-- NOTE: this doesn't constant folding. +run :: MonadError Error m => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.ConstantPropagation" $ do + precondition $ do + ensureWellTyped prog + prog <- return $ run' prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/ConvexHullTrick.hs b/src/Jikka/Core/Convert/ConvexHullTrick.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/ConvexHullTrick.hs @@ -0,0 +1,208 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.ConvexHullTrick +-- Description : uses convex hull trick. / convex hull trick を使います。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- \[ +-- \newcommand\int{\mathbf{int}} +-- \newcommand\bool{\mathbf{bool}} +-- \newcommand\list{\mathbf{list}} +-- \] +module Jikka.Core.Convert.ConvexHullTrick + ( run, + + -- * internal rules + rule, + parseLinearFunctionBody, + parseLinearFunctionBody', + ) +where + +import Control.Monad.Trans.Maybe +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Language.ArithmeticalExpr +import Jikka.Core.Language.Beta +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.FreeVars +import Jikka.Core.Language.Lint +import Jikka.Core.Language.RewriteRules +import Jikka.Core.Language.Util + +hoistMaybe :: Applicative m => Maybe a -> MaybeT m a +hoistMaybe = MaybeT . pure + +-- | This is something commutative because only one kind of @c@ is allowed. +plusPair :: (ArithmeticalExpr, ArithmeticalExpr) -> (ArithmeticalExpr, ArithmeticalExpr) -> Maybe (ArithmeticalExpr, ArithmeticalExpr) +plusPair (a1, c1) (a2, _) | isZeroArithmeticalExpr a2 = Just (a1, c1) +plusPair (a1, c1) (_, c2) | isZeroArithmeticalExpr c2 = Just (a1, c1) +plusPair (a1, _) (a2, c2) | isZeroArithmeticalExpr a1 = Just (a2, c2) +plusPair (_, c1) (a2, c2) | isZeroArithmeticalExpr c1 = Just (a2, c2) +plusPair (a1, c1) (a2, c2) = + let (k1, c1') = splitConstantFactorArithmeticalExpr c1 + (k2, c2') = splitConstantFactorArithmeticalExpr c2 + a1' = multArithmeticalExpr (integerArithmeticalExpr k1) a1 + a2' = multArithmeticalExpr (integerArithmeticalExpr k2) a2 + in if c1' == c2' + then Just (plusArithmeticalExpr a1' a2', c1') + else Nothing + +sumPairs :: [(ArithmeticalExpr, ArithmeticalExpr)] -> Maybe (ArithmeticalExpr, ArithmeticalExpr) +sumPairs = foldr (\e1 e2 -> plusPair e1 =<< e2) (Just (integerArithmeticalExpr 1, integerArithmeticalExpr 0)) + +-- | `parseLinearFunctionBody'` parses the body of a linear function which can be decomposed to convex hull trick. +-- @parseLinearFunctionBody' f i j e@ finds a 4-tuple @a, b, c, d@ where @e = a(f[j], j) c(f[< i], i) + b(f[j], j) + d(f[< i], i)@. +-- +-- TODO: What is the relation between @j@ and @k@? +parseLinearFunctionBody' :: VarName -> VarName -> VarName -> Expr -> Maybe (Expr, Expr, Expr, Expr) +parseLinearFunctionBody' f i j e = result <$> go e + where + result (a, c, b, d) = + let (k, a') = splitConstantFactorArithmeticalExpr a + c' = multArithmeticalExpr (integerArithmeticalExpr k) c + in (formatArithmeticalExpr a', formatArithmeticalExpr c', formatArithmeticalExpr b, formatArithmeticalExpr d) + go = \case + Negate' e -> do + (a, c, b, d) <- go e + return (a, negateArithmeticalExpr c, negateArithmeticalExpr b, negateArithmeticalExpr d) + Plus' e1 e2 -> do + (a1, c1, b1, d1) <- go e1 + (a2, c2, b2, d2) <- go e2 + (a, c) <- plusPair (a1, c1) (a2, c2) + return (a, c, plusArithmeticalExpr b1 b2, plusArithmeticalExpr d1 d2) + Minus' e1 e2 -> do + (a1, c1, b1, d1) <- go e1 + (a2, c2, b2, d2) <- go e2 + (a, c) <- plusPair (a1, c1) (negateArithmeticalExpr a2, c2) + return (a, c, minusArithmeticalExpr b1 b2, minusArithmeticalExpr d1 d2) + Mult' e1 e2 -> do + (a1, c1, b1, d1) <- go e1 + (a2, c2, b2, d2) <- go e2 + (a, c) <- + sumPairs + [ (multArithmeticalExpr a1 a2, multArithmeticalExpr c1 c2), + (multArithmeticalExpr b2 a1, c1), + (multArithmeticalExpr b1 a2, c2), + (a1, multArithmeticalExpr c1 d2), + (a2, multArithmeticalExpr c2 d1), + (b2, d1), + (b1, d2) + ] + return (a, c, multArithmeticalExpr b1 b2, multArithmeticalExpr d1 d2) + e + | f `isUnusedVar` e && j `isUnusedVar` e -> + -- NOTE: Put constants to @d@ and simplify @a, b@ + return (integerArithmeticalExpr 1, integerArithmeticalExpr 0, integerArithmeticalExpr 0, parseArithmeticalExpr e) + e + | f `isUnusedVar` e && i `isUnusedVar` e -> + return (integerArithmeticalExpr 1, integerArithmeticalExpr 0, parseArithmeticalExpr e, integerArithmeticalExpr 0) + e@(At' _ (Var f') index) | f' == f -> case unNPlusKPattern (parseArithmeticalExpr index) of + Just (i', k) | i' == i && k < 0 -> do + return (integerArithmeticalExpr 1, integerArithmeticalExpr 0, integerArithmeticalExpr 0, parseArithmeticalExpr e) + Just (j', 0) | j' == j -> do + return (integerArithmeticalExpr 1, integerArithmeticalExpr 0, parseArithmeticalExpr e, integerArithmeticalExpr 0) + _ -> Nothing + _ -> Nothing + +parseLinearFunctionBody :: MonadAlpha m => VarName -> VarName -> Integer -> Expr -> m (Maybe (Expr, Expr, Expr, Expr, Expr)) +parseLinearFunctionBody f i k = runMaybeT . go + where + go = \case + Min1' _ (Map' _ _ (Lam j _ step) (Range1' size)) -> case unNPlusKPattern (parseArithmeticalExpr size) of + Just (i', k') | i' == i && k' == k -> do + (a, b, c, d) <- hoistMaybe $ parseLinearFunctionBody' f i j step + -- raname @j@ to @i@ + a <- lift $ substitute j (Var i) a + c <- lift $ substitute j (Var i) c + return (LitInt' 1, a, b, c, d) + _ -> hoistMaybe Nothing + Max1' _ (Map' _ _ (Lam j _ step) (Range1' size)) -> case unNPlusKPattern (parseArithmeticalExpr size) of + Just (i', k') | i' == i && k' == k -> do + (a, b, c, d) <- hoistMaybe $ parseLinearFunctionBody' f i j step + -- raname @j@ to @i@ + a <- lift $ substitute j (Var i) a + c <- lift $ substitute j (Var i) c + return (LitInt' (-1), a, Negate' b, Negate' c, d) + _ -> hoistMaybe Nothing + Negate' e -> do + (sign, a, b, c, d) <- go e + return (Negate' sign, a, b, c, Negate' d) + Plus' e1 e2 | isConstantTimeExpr e2 -> do + (sign, a, b, c, d) <- go e1 + return (sign, a, b, c, Plus' d e2) + Plus' e1 e2 | isConstantTimeExpr e1 -> do + (sign, a, b, c, d) <- go e2 + return (sign, a, b, c, Plus' e1 d) + Minus' e1 e2 | isConstantTimeExpr e2 -> do + (sign, a, b, c, d) <- go e1 + return (sign, a, b, c, Minus' d e2) + Minus' e1 e2 | isConstantTimeExpr e1 -> do + (sign, a, b, c, d) <- go e2 + return (Negate' sign, a, b, c, Minus' e1 d) + Mult' e1 e2 | isConstantTimeExpr e2 -> do + (sign, a, b, c, d) <- go e1 + return (Mult' sign e2, a, b, c, Mult' d e2) + Mult' e1 e2 | isConstantTimeExpr e1 -> do + (sign, a, b, c, d) <- go e2 + return (Mult' e1 sign, a, b, c, Mult' e1 d) + _ -> hoistMaybe Nothing + +getLength :: Expr -> Maybe Integer +getLength = \case + Nil' _ -> Just 0 + Cons' _ _ xs -> succ <$> getLength xs + Snoc' _ xs _ -> succ <$> getLength xs + _ -> Nothing + +rule :: (MonadAlpha m, MonadError Error m) => RewriteRule m +rule = RewriteRule $ \_ -> \case + -- build (fun f -> step(f)) base n + Build' IntTy (Lam f _ step) base n -> runMaybeT $ do + i <- lift genVarName' + k <- hoistMaybe $ getLength base + step <- replaceLenF f i k step + -- step(f) = sign(f) * min (map (fun j -> a(f, j) c(f) + b(f, j)) (range (i + k))) + d(f) + (sign, a, c, b, d) <- MaybeT $ parseLinearFunctionBody f i k step + x <- lift genVarName' + y <- lift genVarName' + f' <- lift $ genVarName f + let ts = [ConvexHullTrickTy, ListTy IntTy] + -- base' = (empty, base) + let base' = uncurryApp (Tuple' ts) [ConvexHullTrickInit', base] + -- step' = fun (cht, f) i -> + -- let f' = setat f index(i) (min cht f[i + k] + c(i)) + -- in let cht' = update cht a(i) b(i) + -- in (cht', f') + let step' = + Lam2 x (TupleTy ts) i IntTy $ + Let f (ListTy IntTy) (Proj' ts 1 (Var x)) $ + Let y ConvexHullTrickTy (ConvexHullTrickInsert' (Proj' ts 0 (Var x)) a b) $ + Let f' (ListTy IntTy) (Snoc' IntTy (Var f) (Plus' (Mult' sign (ConvexHullTrickGetMin' (Var y) c)) d)) $ + uncurryApp (Tuple' ts) [Var y, Var f'] + -- proj 1 (foldl step' base' (range (n - 1))) + return $ Proj' ts 1 (Foldl' IntTy (TupleTy ts) step' base' (Range1' n)) + _ -> return Nothing + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram = applyRewriteRuleProgram' rule + +-- | `run` optimizes a DP which has the recurrence relation +-- \[ +-- \mathrm{dp}(i) = \min a(j) x(i) + b(j) \lbrace \mid j \lt i \rbrace + c(i) +-- \] where only appropriate elements of \(\mathrm{dp}\) are used in \(a, x, b, c\). +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.ConvexHullTrick" $ do + precondition $ do + ensureWellTyped prog + prog <- runProgram prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/CumulativeSum.hs b/src/Jikka/Core/Convert/CumulativeSum.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/CumulativeSum.hs @@ -0,0 +1,86 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.CumulativeSum +-- Description : processes queries like range sum query using cumulative sums. / 累積和を用いて range sum query のようなクエリを処理します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Convert.CumulativeSum + ( run, + + -- * internal rules + rule, + ) +where + +import Data.Maybe +import Jikka.Common.Alpha +import Jikka.Common.Error +import qualified Jikka.Core.Convert.Alpha as Alpha +import Jikka.Core.Language.ArithmeticalExpr +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.FreeVars +import Jikka.Core.Language.Lint +import Jikka.Core.Language.RewriteRules +import Jikka.Core.Language.Util + +cumulativeMax :: MonadAlpha m => (Expr -> Expr -> Expr) -> Type -> Maybe Expr -> Expr -> Expr -> m Expr +cumulativeMax max2 t a0 a n = do + b <- genVarName' + let e = At' t (Var b) n + x1 <- genVarName' + x2 <- genVarName' + let a0' = fromMaybe (At' t a (LitInt' 0)) a0 + return $ Let b (ListTy t) (Scanl' t t (Lam2 x1 t x2 t (max2 (Var x1) (Var x2))) a0' a) e + +rule :: MonadAlpha m => RewriteRule m +rule = RewriteRule $ \_ -> \case + Sum' (Map' _ _ (Lam x _ (At' _ a index)) (Range1' n)) | x `isUnusedVar` a -> do + case makeAffineFunctionFromArithmeticalExpr x (parseArithmeticalExpr index) of + Just (coeff, shift) | isOneArithmeticalExpr coeff -> do + b <- genVarName' + let e = + if isZeroArithmeticalExpr shift + then At' IntTy (Var b) n + else Minus' (At' IntTy (Var b) (Plus' n (formatArithmeticalExpr shift))) (At' IntTy (Var b) (formatArithmeticalExpr shift)) + return . Just $ + Let b (ListTy IntTy) (Scanl' IntTy IntTy (Lit (LitBuiltin Plus)) Lit0 a) e + _ -> return Nothing + Max1' t (Cons' _ a0 (Map' _ _ (Lam x _ (At' _ a (Var x'))) (Range1' n))) | x' == x && x `isUnusedVar` a -> do + Just <$> cumulativeMax (Max2' t) t (Just a0) a n + Max1' t (Map' _ _ (Lam x _ (At' _ a (Var x'))) (Range1' n)) | x' == x && x `isUnusedVar` a -> do + Just <$> cumulativeMax (Max2' t) t Nothing a n + Min1' t (Cons' _ a0 (Map' _ _ (Lam x _ (At' _ a (Var x'))) (Range1' n))) | x' == x && x `isUnusedVar` a -> do + Just <$> cumulativeMax (Min2' t) t (Just a0) a n + Min1' t (Map' _ _ (Lam x _ (At' _ a (Var x'))) (Range1' n)) | x' == x && x `isUnusedVar` a -> do + Just <$> cumulativeMax (Min2' t) t Nothing a n + _ -> return Nothing + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram = applyRewriteRuleProgram' rule + +-- | `run` introduces cumulative sums. +-- +-- == Examples +-- +-- Before: +-- +-- > sum (fun i -> a[i]) (range n) +-- +-- After: +-- +-- > let b = scanl (+) 0 a in b[n] +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.CumulativeSum" $ do + precondition $ do + ensureWellTyped prog + prog <- Alpha.run prog + prog <- runProgram prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/Eta.hs b/src/Jikka/Core/Convert/Eta.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/Eta.hs @@ -0,0 +1,83 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.Eta +-- Description : does eta-reductions and makes exprs pointful. / eta 簡約を行って式を pointful にします。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Convert.Eta + ( run, + + -- * internal rules + rule, + ) +where + +import Data.Maybe +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Lint +import Jikka.Core.Language.RewriteRules +import Jikka.Core.Language.Util + +expandExpr :: MonadAlpha m => Type -> Expr -> m (Maybe Expr) +expandExpr t e = case (t, e) of + (FunTy t1 t2, Lam x _ body) -> do + body <- expandExpr t2 body + return $ Lam x t1 <$> body + (FunTy t1 t2, e) -> do + x <- genVarName' + let e' = App e (Var x) + e'' <- expandExpr t2 e' + return . Just $ Lam x t1 (fromMaybe e' e'') + _ -> return Nothing + +rule :: MonadAlpha m => RewriteRule m +rule = + let go :: MonadAlpha m => Expr -> Type -> (Expr -> Expr) -> m (Maybe Expr) + go e t f = (f <$>) <$> expandExpr t e + in RewriteRule $ \_ -> \case + Let x t e1 e2 -> go e1 t (\e1 -> Let x t e1 e2) + Iterate' t k f x -> go f (FunTy t t) (\f -> Iterate' t k f x) + Foldl' t1 t2 f init xs -> go f (FunTy t2 (FunTy t1 t1)) (\f -> Foldl' t1 t2 f init xs) + Scanl' t1 t2 f init xs -> go f (FunTy t2 (FunTy t1 t1)) (\f -> Scanl' t1 t2 f init xs) + Build' t f xs n -> go f (FunTy (ListTy t) t) (\f -> Build' t f xs n) + Map' t1 t2 f xs -> go f (FunTy t1 t2) (\f -> Map' t1 t2 f xs) + Filter' t f xs -> go f (FunTy t BoolTy) (\f -> Filter' t f xs) + _ -> return Nothing + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram = applyRewriteRuleProgram' rule + +-- `run` does eta-reductions in some locations. +-- This aims to: + +-- * simplify other rewrite-rules + +-- * convert to C++ + +-- TODO: expand in toplevel-let too. +-- +-- == Examples +-- +-- Before: +-- +-- > foldl (+) 0 xs +-- +-- After: +-- +-- > foldl (fun y x -> y + x) 0 xs +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.Eta" $ do + precondition $ do + ensureWellTyped prog + prog <- runProgram prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/MakeScanl.hs b/src/Jikka/Core/Convert/MakeScanl.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/MakeScanl.hs @@ -0,0 +1,293 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.MakeScanl +-- Description : converts @foldl@ on lists to @scanl@. / リスト上の @foldl@ を @scanl@ に変換します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- \[ +-- \newcommand\int{\mathbf{int}} +-- \newcommand\bool{\mathbf{bool}} +-- \newcommand\list{\mathbf{list}} +-- \] +module Jikka.Core.Convert.MakeScanl + ( run, + + -- * internal rules + rule, + reduceScanlBuild, + reduceFoldlSetAtRecurrence, + reduceFoldlSetAtAccumulation, + reduceFoldlSetAtGeneric, + getRecurrenceFormulaBase, + getRecurrenceFormulaStep1, + getRecurrenceFormulaStep, + ) +where + +import Control.Monad.Trans.Maybe +import qualified Data.Map as M +import Data.Maybe +import qualified Data.Vector as V +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Language.ArithmeticalExpr +import Jikka.Core.Language.Beta +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.FreeVars +import Jikka.Core.Language.Lint +import Jikka.Core.Language.RewriteRules +import Jikka.Core.Language.Util + +-- | +-- == List of builtin functions which are reduced +-- +-- * `Nil` \(: \forall \alpha. \list(\alpha)\) +-- * `Cons` \(: \forall \alpha. \alpha \to \list(\alpha) \to \list(\alpha)\) +-- * `Scanl` \(: \forall \alpha \beta. (\beta \to \alpha \to \beta) \to \beta \to \list(\alpha) \to \list(\beta)\) +reduceScanlBuild :: Monad m => RewriteRule m +reduceScanlBuild = simpleRewriteRule $ \case + Scanl' _ t2 _ init (Nil' _) -> Just $ Cons' t2 init (Nil' t2) + Scanl' t1 t2 f init (Cons' _ x xs) -> Just $ Cons' t2 init (Scanl' t1 t2 f (App2 f init x) xs) + _ -> Nothing + +-- | `getRecurrenceFormulaBase` makes a pair @((a_0, ..., a_{k - 1}), a)@ from @setat (... (setat a 0 a_0) ...) (k - 1) a_{k - 1})@. +getRecurrenceFormulaBase :: Expr -> ([Expr], Expr) +getRecurrenceFormulaBase = go (V.replicate recurrenceLimit Nothing) + where + recurrenceLimit :: Num a => a + recurrenceLimit = 20 + go :: V.Vector (Maybe Expr) -> Expr -> ([Expr], Expr) + go base = \case + SetAt' _ e (LitInt' i) e' | 0 <= i && i < recurrenceLimit -> go (base V.// [(fromInteger i, Just e')]) e + e -> (map fromJust (takeWhile isJust (V.toList base)), e) + +-- | `getRecurrenceFormulaStep1` removes `At` in @body@. +getRecurrenceFormulaStep1 :: MonadAlpha m => Int -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr) +getRecurrenceFormulaStep1 shift t a i body = do + x <- genVarName a + let proj k = + if toInteger shift + k == 0 + then Just $ Var x + else Nothing + let go :: Expr -> Maybe Expr + go = \case + At' _ (Var a') i' | a' == a -> case unNPlusKPattern (parseArithmeticalExpr i') of + Just (i', k) | i' == i -> proj k + _ -> Nothing + Var x -> if x == a then Nothing else Just (Var x) + Lit lit -> Just $ Lit lit + App f e -> App <$> go f <*> go e + Lam x t e -> Lam x t <$> if x == a then Just e else go e + Let x t e1 e2 -> Let x t <$> go e1 <*> if x == a then Just e2 else go e2 + return $ case go body of + Just body -> Just $ Lam2 x t i IntTy body + Nothing -> Nothing + +-- | `getRecurrenceFormulaStep` replaces `At` in @body@ with `Proj`. +getRecurrenceFormulaStep :: MonadAlpha m => Int -> Int -> Type -> VarName -> VarName -> Expr -> m (Maybe Expr) +getRecurrenceFormulaStep shift size t a i body = do + x <- genVarName a + let ts = replicate size t + let proj k = + if 0 <= toInteger shift + k && toInteger shift + k < toInteger size + then Just $ Proj' ts (shift + fromInteger k) (Var x) + else Nothing + let go :: Expr -> Maybe Expr + go = \case + At' _ (Var a') i' | a' == a -> case unNPlusKPattern (parseArithmeticalExpr i') of + Just (i', k) | i' == i -> proj k + _ -> Nothing + Var x -> if x == a then Nothing else Just (Var x) + Lit lit -> Just $ Lit lit + App f e -> App <$> go f <*> go e + Lam x t e -> Lam x t <$> if x == a then Just e else go e + Let x t e1 e2 -> Let x t <$> go e1 <*> if x == a then Just e2 else go e2 + return $ case go body of + Just body -> Just $ Lam2 x (TupleTy ts) i IntTy (uncurryApp (Tuple' ts) (map (\i -> Proj' ts i (Var x)) [1 .. size - 1] ++ [body])) + Nothing -> Nothing + +hoistMaybe :: Applicative m => Maybe a -> MaybeT m a +hoistMaybe = MaybeT . pure + +-- | +-- * This assumes that `Range2` and `Range3` are already converted to `Range1` (`Jikka.Core.Convert.ShortCutFusion`). +-- * This assumes that combinations `Foldl` and `Map` squashed (`Jikka.Core.Convert.ShortCutFusion`). +-- * This assumes that constants are already folded (`Jikka.Core.Convert.ConstantFolding`). +reduceFoldlSetAtRecurrence :: MonadAlpha m => RewriteRule m +reduceFoldlSetAtRecurrence = RewriteRule $ \_ -> \case + -- foldl (fun a i -> setat a index(i) step(a, i)) base indices + Foldl' _ (ListTy t2) (Lam2 a _ i _ (SetAt' _ (Var a') index step)) base indices | a' == a && a `isUnusedVar` index -> runMaybeT $ do + -- index(i) = i + k + k <- hoistMaybe $ case unNPlusKPattern (parseArithmeticalExpr index) of + Just (i', k) | i' == i -> Just k + _ -> Nothing + -- indices = range n + n <- hoistMaybe $ case indices of + Range1' n -> Just n -- We can do this because foldl-map combinations are already reduced. + _ -> Nothing + -- init = setat (k-1) a_{k-1} (... (setat 0 a_0 (range n)) ...) + (base, _) <- return $ getRecurrenceFormulaBase base -- TODO: care about cases when base is longer than indices + case base of + [] -> + if k == 0 && a `isUnusedVar` step + then return $ Map' IntTy t2 (Lam i IntTy step) (Range1' n) + else hoistMaybe Nothing + [base] -> do + step <- MaybeT $ getRecurrenceFormulaStep1 (- 1 + fromInteger k) t2 a i step + return $ Scanl' IntTy t2 step base (Range1' n) + _ -> do + let ts = replicate (length base) t2 + let base' = uncurryApp (Tuple' ts) base + step <- MaybeT $ getRecurrenceFormulaStep (- length base + fromInteger k) (length base) t2 a i step + x <- lift (genVarName a) + return $ foldr (Cons' t2) (Map' (TupleTy ts) t2 (Lam x (TupleTy ts) (Proj' ts (length base - 1) (Var x))) (Scanl' IntTy (TupleTy ts) step base' (Range1' n))) (init base) + _ -> return Nothing + +-- | `checkAccumulationFormulaStep` checks that all `At` in @body@ about @a@ are @At a i@. +checkAccumulationFormulaStep :: VarName -> VarName -> Expr -> Bool +checkAccumulationFormulaStep a i = go + where + go = \case + At' _ (Var a') i' | a' == a -> case i' of + Var i' | i' == i -> True + _ -> False + Var x -> x /= a + Lit _ -> True + App f e -> go f && go e + Lam x _ e -> x == a || go e + Let x _ e1 e2 -> go e1 && (x == a || go e2) + +-- | +-- * This assumes that `Range2` and `Range3` are already converted to `Range1` (`Jikka.Core.Convert.ShortCutFusion`). +-- * This assumes that combinations `Foldl` and `Map` squashed (`Jikka.Core.Convert.ShortCutFusion`). +-- * This assumes that constants are already folded (`Jikka.Core.Convert.ConstantFolding`). +reduceFoldlSetAtAccumulation :: MonadAlpha m => RewriteRule m +reduceFoldlSetAtAccumulation = RewriteRule $ \_ -> \case + -- foldl (fun a i -> setat a index() step(a, i)) base indices + Foldl' _ (ListTy t2) (Lam2 a _ i _ (SetAt' _ (Var a') index step)) base indices | a' == a && a `isUnusedVar` index && i `isUnusedVar` index -> runMaybeT $ do + -- step(a, i) = op (at a index()) step'(a, i) + (accumulate, step) <- hoistMaybe $ case step of + Max2' t (At' _ (Var a') index') step | a' == a && index' == index -> Just (Max1' t, step) + Min2' t (At' _ (Var a') index') step | a' == a && index' == index -> Just (Min1' t, step) + Plus' (At' _ (Var a') index') step | a' == a && index' == index -> Just (Sum', step) + Mult' (At' _ (Var a') index') step | a' == a && index' == index -> Just (Product', step) + ModPlus' (At' _ (Var a') index') step m | a' == a && index' == index && a `isUnusedVar` m && i `isUnusedVar` m -> Just ((`ModSum'` m), step) + ModMult' (At' _ (Var a') index') step m | a' == a && index' == index && a `isUnusedVar` m && i `isUnusedVar` m -> Just ((`ModProduct'` m), step) + _ -> Nothing + -- indices = range (index()) + guard $ indices == Range1' index + -- step'(a, i) = step''(at a i) + guard $ checkAccumulationFormulaStep a i step + step <- lift $ substitute a base step + return $ SetAt' t2 base index (accumulate (Map' IntTy t2 (Lam i IntTy step) (Range1' index))) + _ -> return Nothing + +-- | `checkGenericRecurrenceFormulaStep` checks that all `At` in @body@ about @a@ have indices less than @i + k@. +checkGenericRecurrenceFormulaStep :: VarName -> VarName -> Integer -> Expr -> Bool +checkGenericRecurrenceFormulaStep a = \i k -> go (M.fromList [(i, k - 1)]) + where + -- (i, k) in env menas a[i + k] is accessible but a[i + k + 1] is not. + go :: M.Map VarName Integer -> Expr -> Bool + go env = \case + At' _ (Var a') i | a' == a -> case unNPlusKPattern (parseArithmeticalExpr i) of + Just (i, k) -> case M.lookup i env of + Just limit -> k <= limit + Nothing -> False + _ -> False + Map' _ _ (Lam j _ body) (Range1' n) | j /= a -> case unNPlusKPattern (parseArithmeticalExpr n) of + Just (i, k) -> case M.lookup i env of + Just limit -> go (M.insert j (limit - k + 1) env) body + Nothing -> go env body && go env n + _ -> go env body && go env n + Var x -> x /= a + Lit _ -> True + App f e -> go env f && go env e + Lam x _ e -> x == a || go env e + Let x _ e1 e2 -> go env e1 && (x == a || go env e2) + +reduceFoldlSetAtGeneric :: MonadAlpha m => RewriteRule m +reduceFoldlSetAtGeneric = RewriteRule $ \_ -> \case + -- foldl (fun a i -> setat a index(i) step(a, i)) base indices + Foldl' _ (ListTy t2) (Lam2 a _ i _ (SetAt' _ (Var a') index step)) base indices | a' == a && a `isUnusedVar` index -> runMaybeT $ do + -- index(i) = i + k + k <- hoistMaybe $ case unNPlusKPattern (parseArithmeticalExpr index) of + Just (i', k) | i' == i -> Just k + _ -> Nothing + -- indices = range n + n <- hoistMaybe $ case indices of + Range1' n -> Just n -- We can do this because foldl-map combinations are already reduced. + _ -> Nothing + -- base = setat (k - 1) a_{k - 1} (... (setat 0 a_0 (range n)) ...) + (base, _) <- return $ getRecurrenceFormulaBase base -- TODO: care about cases when base is longer than indices + -- step(a, i) = step(a[0], a[1], ..., a[i + k - 1], i) + guard $ checkGenericRecurrenceFormulaStep a i k step + step <- lift $ substitute i (Minus' (Len' t2 (Var a)) (LitInt' k)) step + return $ Build' t2 (Lam a (ListTy t2) step) (foldl (Snoc' t2) (Nil' t2) base) n + _ -> return Nothing + +rule :: MonadAlpha m => RewriteRule m +rule = + mconcat + [ reduceScanlBuild, + reduceFoldlSetAtRecurrence, + reduceFoldlSetAtAccumulation, + reduceFoldlSetAtGeneric + ] + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram = applyRewriteRuleProgram' rule + +-- | `run` replaces `Foldl` with `Scanl`. +-- +-- == Example +-- +-- Before: +-- +-- > let xs = range n +-- > xs[0] <- 0 +-- > xs[1] <- 1 +-- > foldl (fun a i -> do +-- > xs[i + 2] <- xs[i] + xs[i + 1] +-- > xs +-- > ) xs (range (n - 2)) +-- +-- After: +-- +-- > 0 : map snd ( +-- > scanl (fun a i -> (snd a, fst a + snd a)) +-- > (0, 1) +-- > (range (n - 2))) +-- +-- == List of builtin functions which are reduced +-- +-- === Build functions +-- +-- * `Nil` \(: \forall \alpha. \list(\alpha)\) +-- * `Cons` \(: \forall \alpha. \alpha \to \list(\alpha) \to \list(\alpha)\) +-- * `Range1` \(: \int \to \list(\int)\) +-- * `Build` \(: \forall \alpha. (\list(\alpha) \to \alpha) \to \list(\alpha) \to \int \to \list(\alpha)\) +-- +-- === Map functions +-- +-- * `Scanl` \(: \forall \alpha \beta. (\beta \to \alpha \to \beta) \to \beta \to \list(\alpha) \to \list(\beta)\) +-- * `SetAt` \(: \forall \alpha. \list(\alpha) \to \int \to \alpha \to \list(\alpha)\) +-- +-- === Fold functions +-- +-- * `Foldl` \(: \forall \alpha \beta. (\beta \to \alpha \to \beta) \to \beta \to \list(\alpha) \to \beta\) +-- * `At` \(: \forall \alpha. \list(\alpha) \to \int \to \alpha\) +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.MakeScanl" $ do + precondition $ do + ensureWellTyped prog + prog <- runProgram prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/MatrixExponentiation.hs b/src/Jikka/Core/Convert/MatrixExponentiation.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/MatrixExponentiation.hs @@ -0,0 +1,140 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.MatrixExponentiation +-- Description : replaces repeated applications of linear (or, affine) functions with powers of matrices. / 線形な (あるいは affine な) 関数の繰り返しの適用を行列累乗で置き換えます。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Convert.MatrixExponentiation + ( run, + ) +where + +import Control.Monad.Trans +import Control.Monad.Trans.Maybe +import qualified Data.Vector as V +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Common.Matrix +import Jikka.Core.Language.ArithmeticalExpr +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.FreeVars +import Jikka.Core.Language.Lint +import Jikka.Core.Language.RewriteRules +import Jikka.Core.Language.Util + +toLinearExpression :: VarName -> Expr -> Maybe (Maybe Expr, Maybe Expr) +toLinearExpression x e = do + (a, b) <- makeVectorFromArithmeticalExpr (V.singleton x) (parseArithmeticalExpr e) + case V.toList a of + [a] -> + let a' = if isOneArithmeticalExpr a then Nothing else Just (formatArithmeticalExpr a) + b' = if isZeroArithmeticalExpr b then Nothing else Just (formatArithmeticalExpr b) + in Just (a', b') + _ -> error $ "Jikka.Core.Convert.MatrixExponentiation.toLinearExpression: size mismtach: " ++ show (V.length a) + +fromMatrix :: Matrix ArithmeticalExpr -> Expr +fromMatrix f = + let (h, w) = matsize f + go row = uncurryApp (Tuple' (replicate w IntTy)) (map formatArithmeticalExpr (V.toList row)) + in uncurryApp (Tuple' (replicate h (TupleTy (replicate w IntTy)))) (map go (V.toList (unMatrix f))) + +fromAffineMatrix :: Matrix ArithmeticalExpr -> V.Vector ArithmeticalExpr -> Expr +fromAffineMatrix a b | fst (matsize a) /= V.length b = error $ "Jikka.Core.Convert.MatrixExponentiation.fromAffineMatrix: size mismtach: " ++ show (matsize a) ++ " and " ++ show (V.length b) +fromAffineMatrix a b = + let (h, w) = matsize a + go row c = uncurryApp (Tuple' (replicate (w + 1) IntTy)) (map formatArithmeticalExpr (V.toList row ++ [c])) + bottom = uncurryApp (Tuple' (replicate (w + 1) IntTy)) (replicate w (LitInt' 0) ++ [LitInt' 1]) + in uncurryApp (Tuple' (replicate (h + 1) (TupleTy (replicate (w + 1) IntTy)))) (V.toList (V.zipWith go (unMatrix a) b) ++ [bottom]) + +toMatrix :: MonadAlpha m => [(VarName, Type)] -> VarName -> Int -> Expr -> m (Maybe (Matrix ArithmeticalExpr, Maybe (V.Vector ArithmeticalExpr))) +toMatrix env x n step = + case curryApp step of + (Tuple' _, es) -> runMaybeT $ do + xs <- V.fromList <$> replicateM n (lift (genVarName x)) + let unpackTuple _ e = case e of + Proj' _ i (Var x') | x' == x -> Var (xs V.! i) + _ -> e + rows <- MaybeT . return . forM es $ \e -> do + let e' = mapExpr unpackTuple env e + guard $ x `isUnusedVar` e' + makeVectorFromArithmeticalExpr xs (parseArithmeticalExpr e') + a <- MaybeT . return $ makeMatrix (V.fromList (map fst rows)) + let b = if all (isZeroArithmeticalExpr . snd) rows then Nothing else Just (V.fromList (map snd rows)) + return (a, b) + _ -> return Nothing + +addOneToVector :: Int -> VarName -> Expr +addOneToVector n x = + let ts = replicate n IntTy + in uncurryApp (Tuple' (IntTy : ts)) (map (\i -> Proj' ts i (Var x)) [0 .. n - 1] ++ [LitInt' 1]) + +removeOneFromVector :: Int -> VarName -> Expr +removeOneFromVector n x = + let ts = replicate n IntTy + in uncurryApp (Tuple' ts) (map (\i -> Proj' (IntTy : ts) i (Var x)) [0 .. n - 1]) + +rule :: MonadAlpha m => RewriteRule m +rule = RewriteRule $ \env -> \case + Iterate' IntTy k (Lam x _ step) base -> do + let step' = toLinearExpression x step + return $ case step' of + Nothing -> Nothing + Just (Nothing, Nothing) -> Just base + Just (Nothing, Just b) -> Just $ Plus' base (Mult' k b) + Just (Just a, Nothing) -> Just $ Mult' (Pow' a k) base + Just (Just a, Just b) -> + let a' = Pow' a k + b' = Mult' (FloorDiv' (Minus' (Pow' a k) (LitInt' 1)) (Minus' a (LitInt' 1))) b -- This division has no remainder. + in Just $ Plus' (Mult' a' base) b' + Iterate' (TupleTy ts) k (Lam x _ step) base | isVectorTy' ts -> do + let n = length ts + let go n step base = MatAp' n n (MatPow' n step k) base + step <- toMatrix env x n step + case step of + Nothing -> return Nothing + Just (a, Nothing) -> return . Just $ go n (fromMatrix a) base + Just (a, Just b) -> do + y <- genVarName x + z <- genVarName x + return . Just $ + Let y (TupleTy ts) base $ + Let z (TupleTy (IntTy : ts)) (go (n + 1) (fromAffineMatrix a b) (addOneToVector n y)) $ + removeOneFromVector n z + _ -> return Nothing + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram = applyRewriteRuleProgram' rule + +-- | `run` simplifies an affine functions from vectors to vectors in @iterate@ (`Iterate`) functions. +-- +-- == Examples +-- +-- This makes matrix multiplication. Before: +-- +-- > iterate n (fun xs -> (xs[0] + 2 * xs[1], xs[1])) xs +-- +-- After: +-- +-- > matap (matpow ((1, 2), (0, 1)) n) xs +-- +-- Also this works on integers. Before: +-- +-- > iterate n (fun x -> (2 x + 1)) x +-- +-- After: +-- +-- > (2 ** n) * x + (2 ** n - 1) / (n - 1) +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.MatrixExponentiation" $ do + precondition $ do + ensureWellTyped prog + prog <- runProgram prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/PropagateMod.hs b/src/Jikka/Core/Convert/PropagateMod.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/PropagateMod.hs @@ -0,0 +1,208 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.PropagateMod +-- Description : propagates modulo operations, and replaces integer functions with corresponding functions with modulo. / 剰余演算を伝播させ、整数の関数を対応する modulo 付きの関数で置き換えます。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Convert.PropagateMod + ( run, + ) +where + +import Data.Maybe +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Format (formatType) +import Jikka.Core.Language.Beta +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Lint +import Jikka.Core.Language.RewriteRules +import Jikka.Core.Language.TypeCheck +import Jikka.Core.Language.Util + +-- | `Mod` is a newtype to avoid mistakes that swapping left and right of mod-op. +newtype Mod = Mod Expr + +isModulo' :: Expr -> Mod -> Bool +isModulo' e (Mod m) = case e of + FloorMod' _ m' -> m' == m + ModNegate' _ m' -> m' == m + ModPlus' _ _ m' -> m' == m + ModMinus' _ _ m' -> m' == m + ModMult' _ _ m' -> m' == m + ModInv' _ m' -> m' == m + ModPow' _ _ m' -> m' == m + VecFloorMod' _ _ m' -> m' == m + MatFloorMod' _ _ _ m' -> m' == m + ModMatAp' _ _ _ _ m' -> m' == m + ModMatAdd' _ _ _ _ m' -> m' == m + ModMatMul' _ _ _ _ _ m' -> m' == m + ModMatPow' _ _ _ m' -> m' == m + ModSum' _ m' -> m' == m + ModProduct' _ m' -> m' == m + LitInt' n -> case m of + LitInt' m -> 0 <= n && n < m + _ -> False + Proj' ts _ e | isVectorTy' ts -> e `isModulo'` Mod m + Proj' ts _ e | isMatrixTy' ts -> e `isModulo'` Mod m + Map' _ _ f _ -> f `isModulo'` Mod m + Lam _ _ body -> body `isModulo'` Mod m + e@(App _ _) -> case curryApp e of + (e@(Lam _ _ _), _) -> e `isModulo'` Mod m + (Tuple' ts, es) | isVectorTy' ts -> all (`isModulo'` Mod m) es + (Tuple' ts, es) | isMatrixTy' ts -> all (`isModulo'` Mod m) es + _ -> False + _ -> False + +isModulo :: Expr -> Expr -> Bool +isModulo e m = e `isModulo'` Mod m + +putFloorMod :: MonadAlpha m => Mod -> Expr -> m (Maybe Expr) +putFloorMod (Mod m) = + let return' = return . Just + in \case + Negate' e -> return' $ ModNegate' e m + Plus' e1 e2 -> return' $ ModPlus' e1 e2 m + Minus' e1 e2 -> return' $ ModMinus' e1 e2 m + Mult' e1 e2 -> return' $ ModMult' e1 e2 m + Pow' e1 e2 -> return' $ ModPow' e1 e2 m + MatAp' h w e1 e2 -> return' $ ModMatAp' h w e1 e2 m + MatAdd' h w e1 e2 -> return' $ ModMatAdd' h w e1 e2 m + MatMul' h n w e1 e2 -> return' $ ModMatMul' h n w e1 e2 m + MatPow' n e1 e2 -> return' $ ModMatPow' n e1 e2 m + Sum' e -> return' $ ModSum' e m + Product' e -> return' $ ModProduct' e m + LitInt' n -> case m of + LitInt' m -> return' $ LitInt' (n `mod` m) + _ -> return Nothing + Proj' ts i e | isVectorTy' ts -> return' $ Proj' ts i (VecFloorMod' (length ts) e m) + Proj' ts i e + | isMatrixTy' ts -> + let (h, w) = fromJust (sizeOfMatrixTy (TupleTy ts)) + in return' $ Proj' ts i (MatFloorMod' h w e m) + Map' t1 t2 f xs -> do + f <- putFloorMod (Mod m) f + case f of + Nothing -> return Nothing + Just f -> return' $ Map' t1 t2 f xs + Lam x t body -> do + -- TODO: rename only if required + y <- genVarName x + body <- substitute x (Var y) body + body <- putFloorMod (Mod m) body + case body of + Nothing -> return Nothing + Just body -> return' $ Lam y t body + e@(App _ _) -> case curryApp e of + (f@(Lam _ _ _), args) -> do + f <- putFloorMod (Mod m) f + case f of + Nothing -> return Nothing + Just f -> return' $ uncurryApp f args + (Tuple' ts, es) | isVectorTy' ts -> do + es' <- mapM (putFloorMod (Mod m)) es + if all isNothing es' + then return Nothing + else return' $ uncurryApp (Tuple' ts) (zipWith fromMaybe es es') + (Tuple' ts, es) | isMatrixTy (TupleTy ts) -> do + es' <- mapM (putFloorMod (Mod m)) es + if all isNothing es' + then return Nothing + else return' $ uncurryApp (Tuple' ts) (zipWith fromMaybe es es') + _ -> return Nothing + _ -> return Nothing + +putFloorModGeneric :: MonadAlpha m => (Expr -> Mod -> m Expr) -> Mod -> Expr -> m Expr +putFloorModGeneric fallback m e = + if e `isModulo'` m + then return e + else do + e' <- putFloorMod m e + case e' of + Just e' -> return e' + Nothing -> fallback e m + +putFloorModInt :: MonadAlpha m => Mod -> Expr -> m Expr +putFloorModInt = putFloorModGeneric (\e (Mod m) -> return $ FloorMod' e m) + +putMapFloorMod :: MonadAlpha m => Mod -> Expr -> m Expr +putMapFloorMod = putFloorModGeneric fallback + where + fallback e (Mod m) = do + x <- genVarName' + return $ Map' IntTy IntTy (Lam x IntTy (FloorMod' (Var x) m)) e + +putVecFloorMod :: (MonadError Error m, MonadAlpha m) => [(VarName, Type)] -> Mod -> Expr -> m Expr +putVecFloorMod env = putFloorModGeneric fallback + where + fallback e (Mod m) = do + t <- typecheckExpr env e + case t of + TupleTy ts -> return $ VecFloorMod' (length ts) e m + _ -> throwInternalError $ "not a vector: " ++ formatType t + +putMatFloorMod :: (MonadError Error m, MonadAlpha m) => [(VarName, Type)] -> Mod -> Expr -> m Expr +putMatFloorMod env = putFloorModGeneric fallback + where + fallback e (Mod m) = do + t <- typecheckExpr env e + case t of + TupleTy ts@(TupleTy ts' : _) -> return $ MatFloorMod' (length ts) (length ts') e m + _ -> throwInternalError $ "not a matrix: " ++ formatType t + +rule :: (MonadAlpha m, MonadError Error m) => RewriteRule m +rule = + let go1 m f (t1, e1) = Just <$> (f <$> t1 (Mod m) e1 <*> pure m) + go2 m f (t1, e1) (t2, e2) = Just <$> (f <$> t1 (Mod m) e1 <*> t2 (Mod m) e2 <*> pure m) + in RewriteRule $ \env -> \case + ModNegate' e m | not (e `isModulo` m) -> go1 m ModNegate' (putFloorModInt, e) + ModPlus' e1 e2 m | not (e1 `isModulo` m) || not (e2 `isModulo` m) -> go2 m ModPlus' (putFloorModInt, e1) (putFloorModInt, e2) + ModMinus' e1 e2 m | not (e1 `isModulo` m) || not (e2 `isModulo` m) -> go2 m ModMinus' (putFloorModInt, e1) (putFloorModInt, e2) + ModMult' e1 e2 m | not (e1 `isModulo` m) || not (e2 `isModulo` m) -> go2 m ModMult' (putFloorModInt, e1) (putFloorModInt, e2) + ModInv' e m | not (e `isModulo` m) -> go1 m ModInv' (putFloorModInt, e) + ModPow' e1 e2 m | not (e1 `isModulo` m) -> go2 m ModPow' (putFloorModInt, e1) (\_ e -> return e, e2) + ModMatAp' h w e1 e2 m | not (e1 `isModulo` m) || not (e2 `isModulo` m) -> go2 m (ModMatAp' h w) (putMatFloorMod env, e1) (putVecFloorMod env, e2) + ModMatAdd' h w e1 e2 m | not (e1 `isModulo` m) || not (e2 `isModulo` m) -> go2 m (ModMatAdd' h w) (putMatFloorMod env, e1) (putMatFloorMod env, e2) + ModMatMul' h n w e1 e2 m | not (e1 `isModulo` m) || not (e2 `isModulo` m) -> go2 m (ModMatMul' h n w) (putMatFloorMod env, e1) (putMatFloorMod env, e2) + ModMatPow' n e1 e2 m | not (e1 `isModulo` m) -> go2 m (ModMatPow' n) (putMatFloorMod env, e1) (\_ e -> return e, e2) + ModSum' e m | not (e `isModulo` m) -> go1 m ModSum' (putMapFloorMod, e) + ModProduct' e m | not (e `isModulo` m) -> go1 m ModProduct' (putMapFloorMod, e) + FloorMod' e m -> + if e `isModulo` m + then return $ Just e + else putFloorMod (Mod m) e + VecFloorMod' _ e m -> + if e `isModulo` m + then return $ Just e + else putFloorMod (Mod m) e + MatFloorMod' _ _ e m -> + if e `isModulo` m + then return $ Just e + else putFloorMod (Mod m) e + _ -> return Nothing + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram = applyRewriteRuleProgram' rule + +-- | `run` propagates `FloorMod` to leaves of exprs. +-- For example, this converts the following: +-- +-- > mod ((fun x -> x * x + x) y) 1000000007 +-- +-- to: +-- +-- > (fun x -> mod (mod (x * x) 1000000007 + x) 1000000007) y +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.PropagateMod" $ do + precondition $ do + ensureWellTyped prog + prog <- runProgram prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/RemoveUnusedVars.hs b/src/Jikka/Core/Convert/RemoveUnusedVars.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/RemoveUnusedVars.hs @@ -0,0 +1,77 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.RemoveUnusedVars +-- Description : removes unused variables. / 使われていない変数を削除します。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- `Jikka.Language.Core.RemoveUnusedVars` remove unused variables from exprs. +module Jikka.Core.Convert.RemoveUnusedVars + ( run, + run', + ) +where + +import Jikka.Common.Error +import Jikka.Core.Language.Expr +import Jikka.Core.Language.FreeVars (isUnusedVar) +import Jikka.Core.Language.Lint +import Jikka.Core.Language.Util + +runLet :: VarName -> Type -> Expr -> Expr -> Expr +runLet x t e1 e2 + | isUnusedVar x e2 = e2 + | otherwise = Let x t e1 e2 + +runExpr :: Expr -> Expr +runExpr = \case + Var x -> Var x + Lit lit -> Lit lit + App f e -> App (runExpr f) (runExpr e) + Lam x t e -> Lam x t (runExpr e) + Let x t e1 e2 -> runLet x t (runExpr e1) (runExpr e2) + +runToplevelExpr :: ToplevelExpr -> ToplevelExpr +runToplevelExpr = \case + ResultExpr e -> ResultExpr $ runExpr e + ToplevelLet x t e cont -> ToplevelLet x t (runExpr e) (runToplevelExpr cont) + ToplevelLetRec f args ret body cont -> + let body' = runExpr body + cont' = runToplevelExpr cont + in if isUnusedVar f body' + then ToplevelLet f (curryFunTy (map snd args) ret) (curryLam args body') cont' + else ToplevelLetRec f args ret body' cont' + +run' :: Program -> Program +run' = runToplevelExpr + +-- | `run` removes unused variables in given programs. +-- +-- This also removes variables for recursion, i.e. "rec" flags. +-- `ToplevelLetRec` may becomes `ToplevelLet`. +-- +-- For example, this converts +-- +-- > let rec solve x = +-- > let y = 0 +-- > in x +-- > in solve +-- +-- to +-- +-- > let solve x = +-- > x +-- > in solve +run :: MonadError Error m => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.RemoveUnusedVars" $ do + precondition $ do + ensureWellTyped prog + prog <- return $ run' prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/SegmentTree.hs b/src/Jikka/Core/Convert/SegmentTree.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/SegmentTree.hs @@ -0,0 +1,170 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE ViewPatterns #-} + +-- | +-- Module : Jikka.Core.Convert.SegmentTree +-- Description : uses segment trees. / セグメント木を利用します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- \[ +-- \newcommand\int{\mathbf{int}} +-- \newcommand\bool{\mathbf{bool}} +-- \newcommand\list{\mathbf{list}} +-- \] +module Jikka.Core.Convert.SegmentTree + ( run, + + -- * internal rules + rule, + reduceCumulativeSum, + reduceMin, + ) +where + +import Control.Arrow +import Control.Monad.Trans.Maybe +import Data.List +import qualified Data.Map as M +import Data.Maybe +import Jikka.Common.Alpha +import Jikka.Common.Error +import qualified Jikka.Core.Convert.Alpha as Alpha +import Jikka.Core.Language.Beta +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.FreeVars +import Jikka.Core.Language.Lint +import Jikka.Core.Language.RewriteRules +import Jikka.Core.Language.Util + +pattern CumulativeSum t f e es <- + ( \case + Scanl' t t' (Lam2 x1 t'' x2 t''' (App (App f (Var x1')) (Var x2'))) e es + | t == t' && t' == t'' && t'' == t''' && x1 == x1' && x1 `isUnusedVar` f && x2 == x2' && x2 `isUnusedVar` f -> Just (t, f, e, es) + _ -> Nothing -> + Just (t, f, e, es) + ) + where + CumulativeSum t f e es = + let x1 = findUnusedVarName (VarName "y") f + x2 = findUnusedVarName (VarName "x") f + in Scanl' t t (Lam2 x1 t x2 t (App (App f (Var x1)) (Var x2))) e es + +pattern CumulativeSumFlip t f e es <- + ( \case + Scanl' t t' (Lam2 x1 t'' x2 t''' (App (App f (Var x2')) (Var x1'))) e es + | t == t' && t' == t'' && t'' == t''' && x2 == x2' && x2 `isUnusedVar` f && x1 == x1' && x1 `isUnusedVar` f -> Just (t, f, e, es) + _ -> Nothing -> + Just (t, f, e, es) + ) + where + CumulativeSumFlip t f e es = + let x1 = findUnusedVarName (VarName "y") f + x2 = findUnusedVarName (VarName "x") f + in Scanl' t t (Lam2 x1 t x2 t (App (App f (Var x2)) (Var x1))) e es + +builtinToSemigroup :: Builtin -> Maybe Semigroup' +builtinToSemigroup = \case + Plus -> Just SemigroupIntPlus + Min2 IntTy -> Just SemigroupIntMin + Max2 IntTy -> Just SemigroupIntMax + _ -> Nothing + +semigroupToBuiltin :: Semigroup' -> Builtin +semigroupToBuiltin = \case + SemigroupIntPlus -> Plus + SemigroupIntMin -> Min2 IntTy + SemigroupIntMax -> Max2 IntTy + +unCumulativeSum :: Expr -> Expr -> Maybe (Semigroup', Expr) +unCumulativeSum a = \case + CumulativeSum _ (Lit (LitBuiltin op)) b a' | a' == a -> case builtinToSemigroup op of + Just semigrp -> Just (semigrp, b) + Nothing -> Nothing + -- Semigroups must be commutative to use CumulativeSumFlip. + CumulativeSumFlip _ (Lit (LitBuiltin op)) b a' | a' == a -> case builtinToSemigroup op of + Just semigrp -> Just (semigrp, b) + Nothing -> Nothing + _ -> Nothing + +listCumulativeSum :: Expr -> Expr -> [(Semigroup', Expr)] +listCumulativeSum a = mapMaybe (unCumulativeSum a) . listSubExprs + +replaceWithSegtrees :: VarName -> [(Semigroup', Expr)] -> Expr -> Expr +replaceWithSegtrees a segtrees = go M.empty + where + go :: M.Map VarName (Expr, Expr, Semigroup') -> Expr -> Expr + go env = \case + At' _ (check env -> Just (e, b, semigrp)) i -> + let e' = SegmentTreeGetRange' semigrp e (LitInt' 0) i + in AppBuiltin2 (semigroupToBuiltin semigrp) b e' + Var x -> Var x + Lit lit -> Lit lit + App e1 e2 -> App (go env e1) (go env e2) + Lam x t e -> Lam x t $ go (M.delete x env) e + Let x t e1 e2 -> + let e1' = go env e1 + in case check env e1' of + Just (e1', b, semigrp) -> go (M.insert x (e1', b, semigrp) env) e2 + Nothing -> Let x t (go env e1) (go env e2) + check env = \case + Var x -> M.lookup x env + CumulativeSum _ (Lit (LitBuiltin op)) b (Var a') | a' == a -> case lookup op (map (first semigroupToBuiltin) segtrees) of + Just e -> Just (e, b, fromJust (builtinToSemigroup op)) + Nothing -> Nothing + _ -> Nothing + +-- | `reduceCumulativeSum` converts combinations of cumulative sums and array assignments to segment trees. +reduceCumulativeSum :: (MonadAlpha m, MonadError Error m) => RewriteRule m +reduceCumulativeSum = RewriteRule $ \_ -> \case + -- foldl (fun a i -> setat a index(i) e(a, i)) base incides + Foldl' t1 t2 (Lam2 a _ i _ (SetAt' t (Var a') index e)) base indices | a' == a && a `isUnusedVar` index -> runMaybeT $ do + let sums = listCumulativeSum (Var a) e + guard $ not (null sums) + let semigrps = nub (sort (map fst sums)) + let ts = t2 : map SegmentTreeTy semigrps + c <- lift $ genVarName a + let proj i = Proj' ts i (Var c) + let e' = replaceWithSegtrees a (zip semigrps (map proj [1 ..])) e + guard $ e' /= e + e' <- lift $ substitute a (proj 0) e' + b' <- lift $ genVarName a + let updateSegtrees i semigrp = SegmentTreeSetPoint' semigrp (proj i) index (At' t (Var b') index) + let step = Lam2 c (TupleTy ts) i t1 (Let b' t2 (SetAt' t (proj 0) index e') (uncurryApp (Tuple' ts) (Var b' : zipWith updateSegtrees [1 ..] semigrps))) + b <- lift $ genVarName a + let base' = Var b : map (\semigrp -> SegmentTreeInitList' semigrp (Var b)) semigrps + return $ Let b t2 base (Proj' ts 0 (Foldl' t1 (TupleTy ts) step (uncurryApp (Tuple' ts) base') indices)) + _ -> return Nothing + +-- | `reduceFromMin` uses segment trees from accumulation of min/max which are not reducible to cumulative sums. +-- +-- TODO: implement this +reduceMin :: MonadAlpha m => RewriteRule m +reduceMin = RewriteRule $ \_ -> \case + _ -> return Nothing + +rule :: (MonadAlpha m, MonadError Error m) => RewriteRule m +rule = + mconcat + [ reduceCumulativeSum, + reduceMin + ] + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram = applyRewriteRuleProgram' rule + +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.SegmentTree" $ do + precondition $ do + ensureWellTyped prog + prog <- runProgram prog + prog <- Alpha.run prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/ShortCutFusion.hs b/src/Jikka/Core/Convert/ShortCutFusion.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/ShortCutFusion.hs @@ -0,0 +1,241 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.ShortCutFusion +-- Description : does short cut fusion. / short cut fusion を行います。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- \[ +-- \newcommand\int{\mathbf{int}} +-- \newcommand\bool{\mathbf{bool}} +-- \newcommand\list{\mathbf{list}} +-- \] +module Jikka.Core.Convert.ShortCutFusion + ( run, + + -- * internal rules + rule, + reduceBuild, + reduceMapBuild, + reduceMap, + reduceMapMap, + reduceFoldMap, + reduceFold, + reduceFoldBuild, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Format (formatExpr) +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.FreeVars +import Jikka.Core.Language.Lint +import Jikka.Core.Language.RewriteRules +import Jikka.Core.Language.Util + +-- | +-- * `Range1` remains. +-- * `Range2` is removed. +-- * `Range3` is removed. +-- * `Nil` and `Cons` are kept as is. +reduceBuild :: MonadAlpha m => RewriteRule m +reduceBuild = + let return' = return . Just + in RewriteRule $ \_ -> \case + Range2' l r -> do + let n = Minus' r l + x <- genVarName' + let f = Lam x IntTy (Plus' l (Var x)) + return' $ Map' IntTy IntTy f (Range1' n) + Range3' l r step -> do + let n = CeilDiv' (Minus' r l) step + x <- genVarName' + let f = Lam x IntTy (Plus' l (Mult' step (Var x))) + return' $ Map' IntTy IntTy f (Range1' n) + _ -> return Nothing + +reduceMapBuild :: MonadAlpha m => RewriteRule m +reduceMapBuild = + let return' = return . Just + in RewriteRule $ \_ -> \case + -- reduce `Sorted` + Sorted' _ (Nil' t) -> return' $ Nil' t + Sorted' _ (Range1' n) -> return' $ Range1' n + -- reduce `Reversed` + Reversed' _ (Nil' t) -> return' $ Nil' t + Reversed' _ (Range1' n) -> do + x <- genVarName' + let f = Lam x IntTy (Minus' (Minus' n (Var x)) (LitInt' 1)) + return' $ Map' IntTy IntTy f n + -- reduce `Filter` + Filter' _ _ (Nil' t) -> return' $ Nil' t + -- reduce `Map` + Map' _ _ _ (Nil' t) -> return' $ Nil' t + Map' t1 t2 f (Cons' _ x xs) -> return' $ Cons' t2 (App f x) (Map' t1 t2 f xs) + -- others + _ -> return Nothing + +reduceMap :: Monad m => RewriteRule m +reduceMap = + let return' = return . Just + in RewriteRule $ \_ -> \case + -- reduce `Map` + Map' _ _ (LamId _ _) xs -> return' xs + -- reduce `Filter` + Filter' t (Lam _ _ LitFalse) _ -> return' (Nil' t) + Filter' _ (Lam _ _ LitTrue) xs -> return' xs + -- others + _ -> return Nothing + +-- | +-- * Functions are reordered as: +-- * `Sort` and `Reversed` (functions to reorder) are lastly applied to lists +-- * `Map` (functions to modify lists) +-- * `Filter` (funcitons to reduce lengths) is firstly applied to lists +reduceMapMap :: MonadAlpha m => RewriteRule m +reduceMapMap = + let return' = return . Just + in RewriteRule $ \_ -> \case + -- reduce `Map` + Map' _ _ (LamId _ _) xs -> return' xs + Map' _ t3 g (Map' t1 _ f xs) -> do + x <- genVarName' + let h = Lam x t1 (App g (App f (Var x))) + return' $ Map' t1 t3 h xs + Map' t1 t2 f (Reversed' _ xs) -> return' $ Reversed' t2 (Map' t1 t2 f xs) + -- reduce `Filter` + Filter' t2 g (Map' t1 _ f xs) -> do + x <- genVarName' + let h = Lam x t1 (App g (App f (Var x))) + return' $ Map' t1 t2 f (Filter' t1 h xs) + Filter' t g (Filter' _ f xs) -> do + x <- genVarName' + let h = Lam x t (And' (App g (Var x)) (App f (Var x))) + return' $ Filter' t h xs + Filter' t f (Sorted' _ xs) -> return' $ Sorted' t (Filter' t f xs) + Filter' t f (Reversed' _ xs) -> return' $ Reversed' t (Filter' t f xs) + -- reduce `Reversed` + Reversed' _ (Reversed' _ xs) -> return' xs + Reversed' _ (Map' t1 t2 f xs) -> return' $ Map' t1 t2 f (Reversed' t1 xs) + -- reduce `Sorted` + Sorted' t (Reversed' _ xs) -> return' $ Sorted' t xs + Sorted' t (Sorted' _ xs) -> return' $ Sorted' t xs + -- others + _ -> return Nothing + +reduceFoldMap :: MonadAlpha m => RewriteRule m +reduceFoldMap = + let return' = return . Just + in RewriteRule $ \_ -> \case + -- reduce `Reversed` + Len' t (Reversed' _ xs) -> return' $ Len' t xs + Elem' t x (Reversed' _ xs) -> return' $ Elem' t x xs + At' t (Reversed' _ xs) i -> return' $ At' t xs (Minus' (Minus' (Len' t xs) i) Lit1) + -- reduce `Sorted` + Len' t (Sorted' _ xs) -> return' $ Len' t xs + Elem' t x (Sorted' _ xs) -> return' $ Elem' t x xs + -- reduce `Map` + Len' _ (Map' t1 _ _ xs) -> return' $ Len' t1 xs + At' _ (Map' t1 _ f xs) i -> return' $ App f (At' t1 xs i) + Foldl' _ t3 g init (Map' t1 _ f xs) -> do + x3 <- genVarName' + x1 <- genVarName' + return' $ Foldl' t1 t3 (Lam2 x3 t3 x1 t1 (App2 g (Var x3) (App f (Var x1)))) init xs + -- others + _ -> return Nothing + +reduceFold :: Monad m => RewriteRule m +reduceFold = simpleRewriteRule $ \case + Foldl' t1 t2 (Lam2 x2 _ x1 _ body) init xs | x1 `isUnusedVar` body -> Just $ Iterate' t2 (Len' t1 xs) (Lam x2 t2 body) init + _ -> Nothing + +reduceFoldBuild :: MonadAlpha m => RewriteRule m +reduceFoldBuild = + let return' = return . Just + in RewriteRule $ \_ -> \case + -- reduce `Foldl` + Foldl' _ _ _ init (Nil' _) -> return' init + Foldl' t1 t2 g init (Cons' _ x xs) -> return' $ Foldl' t1 t2 g (App2 g init x) xs + -- reduce `Len` + Len' _ (Nil' _) -> return' Lit0 + Len' t (Cons' _ _ xs) -> return' $ Plus' Lit1 (Len' t xs) + Len' _ (Range1' n) -> return' n + -- reduce `At` + At' t (Nil' _) i -> return' $ Bottom' t $ "cannot subscript empty list: index = " ++ formatExpr i + At' t (Cons' _ x xs) i -> return' $ If' t (Equal' IntTy i Lit0) x (At' t xs (Minus' i Lit1)) + At' _ (Range1' _) i -> return' i + -- reduce `Elem` + Elem' _ _ (Nil' _) -> return' LitFalse + Elem' t y (Cons' _ x xs) -> return' $ And' (Equal' t x y) (Elem' t y xs) + Elem' _ x (Range1' n) -> return' $ And' (LessEqual' IntTy Lit0 x) (LessThan' IntTy x n) + -- others + _ -> return Nothing + +rule :: MonadAlpha m => RewriteRule m +rule = + mconcat + [ reduceFoldMap, + reduceMap, + reduceMapMap, + reduceFoldBuild, + reduceMapBuild, + reduceBuild, + reduceFold + ] + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram = applyRewriteRuleProgram' rule + +-- | `run` does short cut fusion. +-- +-- * This function is mainly for polymorphic reductions. This dosn't do much about concrete things, e.g., arithmetical operations. +-- * This doesn't do nothing about `Scanl` or `SetAt`. +-- +-- == Example +-- +-- Before: +-- +-- > length (map f (cons (-1) (range n))) +-- +-- After: +-- +-- > n + 1 +-- +-- == List of builtin functions which are reduced +-- +-- === Build functions +-- +-- * `Nil` \(: \forall \alpha. \list(\alpha)\) +-- * `Cons` \(: \forall \alpha. \alpha \to \list(\alpha) \to \list(\alpha)\) +-- * `Range1` \(: \int \to \list(\int)\) +-- * `Range2` \(: \int \to \int \to \list(\int)\) +-- * `Range3` \(: \int \to \int \to \int \to \list(\int)\) +-- +-- === Map functions +-- +-- * `Map` \(: \forall \alpha \beta. (\alpha \to \beta) \to \list(\alpha) \to \list(\beta)\) +-- * `Filter` \(: \forall \alpha \beta. (\alpha \to \bool) \to \list(\alpha) \to \list(\beta)\) +-- * `Reversed` \(: \forall \alpha. \list(\alpha) \to \list(\alpha)\) +-- * `Sorted` \(: \forall \alpha. \list(\alpha) \to \list(\alpha)\) +-- +-- === Fold functions +-- +-- * `Foldl` \(: \forall \alpha \beta. (\beta \to \alpha \to \beta) \to \beta \to \list(\alpha) \to \beta\) +-- * `Len` \(: \forall \alpha. \list(\alpha) \to \int\) +-- * `At` \(: \forall \alpha. \list(\alpha) \to \int \to \alpha\) +-- * `Elem` \(: \forall \alpha. \alpha \to \list(\alpha) \to \bool\) +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.ShortCutFusion" $ do + precondition $ do + ensureWellTyped prog + prog <- runProgram prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/SpecializeFoldl.hs b/src/Jikka/Core/Convert/SpecializeFoldl.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/SpecializeFoldl.hs @@ -0,0 +1,108 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.SpecializeFoldl +-- Description : specializes @foldl@ with concrete functions like @sum@ and @product@. / @sum@ や @product@ のような具体的な関数で @foldl@ を特殊化します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- \[ +-- \newcommand\int{\mathbf{int}} +-- \newcommand\bool{\mathbf{bool}} +-- \newcommand\list{\mathbf{list}} +-- \] +module Jikka.Core.Convert.SpecializeFoldl + ( run, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.FreeVars +import Jikka.Core.Language.Lint +import Jikka.Core.Language.RewriteRules + +rule :: MonadAlpha m => RewriteRule m +rule = simpleRewriteRule $ \case + Foldl' t1 t2 (Lam2 x2 _ x1 _ body) init xs -> case body of + -- Sum + Plus' (Var x2') e | x2' == x2 && x2 `isUnusedVar` e -> Just $ Sum' (Cons' t2 init (Map' t1 t2 (Lam x1 t1 e) xs)) + Plus' e (Var x2') | x2' == x2 && x2 `isUnusedVar` e -> Just $ Sum' (Cons' t2 init (Map' t1 t2 (Lam x1 t1 e) xs)) + Minus' (Var x2') e | x2' == x2 && x2 `isUnusedVar` e -> Just $ Minus' init (Sum' (Map' t1 t2 (Lam x1 t1 e) xs)) + -- Product + Mult' (Var x2') e | x2' == x2 && x2 `isUnusedVar` e -> Just $ Product' (Cons' t2 init (Map' t1 t2 (Lam x1 t1 e) xs)) + Mult' e (Var x2') | x2' == x2 && x2 `isUnusedVar` e -> Just $ Product' (Cons' t2 init (Map' t1 t2 (Lam x1 t1 e) xs)) + -- All + And' (Var x2') e | x2' == x2 && x2 `isUnusedVar` e -> Just $ All' (Cons' t2 init (Map' t1 t2 (Lam x1 t1 e) xs)) + And' e (Var x2') | x2' == x2 && x2 `isUnusedVar` e -> Just $ All' (Cons' t2 init (Map' t1 t2 (Lam x1 t1 e) xs)) + -- Any + Or' (Var x2') e | x2' == x2 && x2 `isUnusedVar` e -> Just $ Any' (Cons' t2 init (Map' t1 t2 (Lam x1 t1 e) xs)) + Or' e (Var x2') | x2' == x2 && x2 `isUnusedVar` e -> Just $ Any' (Cons' t2 init (Map' t1 t2 (Lam x1 t1 e) xs)) + -- Max1 + Max2' _ (Var x2') e | x2' == x2 && x2 `isUnusedVar` e -> Just $ Max1' t2 (Cons' t2 init (Map' t1 t2 (Lam x1 t1 e) xs)) + Max2' _ e (Var x2') | x2' == x2 && x2 `isUnusedVar` e -> Just $ Max1' t2 (Cons' t2 init (Map' t1 t2 (Lam x1 t1 e) xs)) + -- Max1 + Min2' _ (Var x2') e | x2' == x2 && x2 `isUnusedVar` e -> Just $ Min1' t2 (Cons' t2 init (Map' t1 t2 (Lam x1 t1 e) xs)) + Min2' _ e (Var x2') | x2' == x2 && x2 `isUnusedVar` e -> Just $ Min1' t2 (Cons' t2 init (Map' t1 t2 (Lam x1 t1 e) xs)) + -- others + _ -> Nothing + -- The outer floor-mod is required because foldl for empty lists returns values without modulo. + FloorMod' (Foldl' t1 t2 (Lam2 x2 _ x1 _ body) init xs) m -> case body of + -- ModSum + ModPlus' (Var x2') e m' | x2' == x2 && x2 `isUnusedVar` e && m' == m -> Just $ ModSum' (Cons' t2 init (Map' t1 t2 (Lam x1 t1 e) xs)) m + ModPlus' e (Var x2') m' | x2' == x2 && x2 `isUnusedVar` e && m' == m -> Just $ ModSum' (Cons' t2 init (Map' t1 t2 (Lam x1 t1 e) xs)) m + ModMinus' (Var x2') e m' | x2' == x2 && x2 `isUnusedVar` e && m' == m -> Just $ ModMinus' init (ModSum' (Map' t1 t2 (Lam x1 t1 e) xs) m) m + -- ModProduct + ModMult' (Var x2') e m' | x2' == x2 && x2 `isUnusedVar` e && m' == m -> Just $ ModProduct' (Cons' t2 init (Map' t1 t2 (Lam x1 t1 e) xs)) m + ModMult' e (Var x2') m' | x2' == x2 && x2 `isUnusedVar` e && m' == m -> Just $ ModProduct' (Cons' t2 init (Map' t1 t2 (Lam x1 t1 e) xs)) m + -- others + _ -> Nothing + -- others + _ -> Nothing + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram = applyRewriteRuleProgram' rule + +-- | `run` reduces summations and products. +-- +-- == Example +-- +-- Before: +-- +-- > foldl (fun x y -> x + y) 0 xs +-- +-- After: +-- +-- > sum xs +-- +-- == List of builtin functions which are reduced +-- +-- === Source functions +-- +-- * `Foldl` \(: \forall \alpha \beta. (\beta \to \alpha \to \beta) \to \beta \to \list(\alpha) \to \beta\) +-- +-- === Destination functions +-- +-- * `Sum` \(: \list(\int) \to \int\) +-- * `Product` \(: \list(\int) \to \int\) +-- * `ModSum` \(: \list(\int) \to \int \to \int\) +-- * `ModProduct` \(: \list(\int) \to \int \to \int\) +-- * `All` \(: \list(\bool) \to \bool\) +-- * `Any` \(: \list(\bool) \to \bool\) +-- * `Max1` \(: \forall \alpha. \list(\alpha) \to \alpha\) +-- * `Min1` \(: \forall \alpha. \list(\alpha) \to \alpha\) +-- * `Iterate` \(: \forall \alpha. \int \to (\alpha \to \alpha) \to \alpha \to \alpha\) +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.SpecializeFoldl" $ do + precondition $ do + ensureWellTyped prog + prog <- runProgram prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/StrengthReduction.hs b/src/Jikka/Core/Convert/StrengthReduction.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/StrengthReduction.hs @@ -0,0 +1,107 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.StrengthReduction +-- Description : does strength reduction. / 演算子強度低減を行います。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Convert.StrengthReduction + ( run, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Lint +import Jikka.Core.Language.RewriteRules + +-- | `eliminateSomeBuiltins` removes some `Builtin` from `Expr` at all. +eliminateSomeBuiltins :: Monad m => RewriteRule m +eliminateSomeBuiltins = simpleRewriteRule $ \case + -- advanced arithmetical functions + Abs' e -> Just $ Max2' IntTy e (Negate' e) + Lcm' e1 e2 -> Just $ FloorDiv' (Gcd' e1 e2) (Mult' e1 e2) + -- logical functions + Implies' e1 e2 -> Just $ Or' (Not' e1) e2 + -- comparison + GreaterThan' t e1 e2 -> Just $ LessThan' t e2 e1 + GreaterEqual' t e1 e2 -> Just $ LessEqual' t e2 e1 + NotEqual' t e1 e2 -> Just $ Not' (Equal' t e1 e2) + _ -> Nothing + +-- | `reduceNegate` brings `Negate` to the root. +reduceNegate :: Monad m => RewriteRule m +reduceNegate = simpleRewriteRule $ \case + Negate' (Negate' e) -> Just e + Plus' (Negate' e1) (Negate' e2) -> Just $ Negate' (Plus' e1 e2) + Minus' e1 (Negate' e2) -> Just $ Plus' e1 e2 + Minus' (Negate' e1) e2 -> Just $ Negate' (Minus' e1 e2) + -- `Minus` is already removed. + Mult' (Negate' e1) e2 -> Just $ Negate' (Mult' e1 e2) + Mult' e1 (Negate' e2) -> Just $ Negate' (Mult' e1 e2) + -- `Abs` is already removed. + Min2' IntTy (Negate' e1) (Negate' e2) -> Just $ Negate' (Max2' IntTy e1 e2) + Max2' IntTy (Negate' e1) (Negate' e2) -> Just $ Negate' (Min2' IntTy e1 e2) + _ -> Nothing + +-- | `reduceNot` brings `Not` to the root. +reduceNot :: Monad m => RewriteRule m +reduceNot = simpleRewriteRule $ \case + Not' (Not' e) -> Just e + And' (Not' e1) (Not' e2) -> Just $ Not' (Or' e1 e2) + Or' (Not' e1) (Not' e2) -> Just $ Not' (And' e1 e2) + -- `Implies` is already removed. + Mult' (Negate' e1) e2 -> Just $ Negate' (Mult' e1 e2) + Mult' e1 (Negate' e2) -> Just $ Negate' (Mult' e1 e2) + If' t (Not' e1) e2 e3 -> Just $ If' t e1 e3 e2 + _ -> Nothing + +-- | `reduceBitNot` brings `BitNot` to the root. +reduceBitNot :: Monad m => RewriteRule m +reduceBitNot = simpleRewriteRule $ \case + BitNot' (BitNot' e) -> Just e + BitAnd' (BitNot' e1) (BitNot' e2) -> Just $ BitNot' (BitOr' e1 e2) + BitOr' (BitNot' e1) (BitNot' e2) -> Just $ BitNot' (BitAnd' e1 e2) + BitXor' (BitNot' e1) e2 -> Just $ BitNot' (BitXor' e1 e2) + BitXor' e1 (BitNot' e2) -> Just $ BitNot' (BitXor' e1 e2) + _ -> Nothing + +misc :: Monad m => RewriteRule m +misc = simpleRewriteRule $ \case + -- arithmetical functions + Pow' (Pow' e1 e2) e3 -> Just $ Pow' e1 (Plus' e2 e3) + -- advanced arithmetical functions + Gcd' (Mult' k1 e1) (Mult' k2 e2) | k1 == k2 -> Just $ Mult' k1 (Gcd' e1 e2) + Gcd' (Mult' k1 e1) (Mult' e2 k2) | k1 == k2 -> Just $ Mult' k1 (Gcd' e1 e2) + Gcd' (Mult' e1 k1) (Mult' e2 k2) | k1 == k2 -> Just $ Mult' k1 (Gcd' e1 e2) + Gcd' (Mult' e1 k1) (Mult' k2 e2) | k1 == k2 -> Just $ Mult' k1 (Gcd' e1 e2) + _ -> Nothing + +rule :: MonadAlpha m => RewriteRule m +rule = + mconcat + [ eliminateSomeBuiltins, + reduceNegate, + reduceNot, + reduceBitNot, + misc + ] + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram = applyRewriteRuleProgram' rule + +-- | TODO: Split and remove this module. +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.StrengthReduction" $ do + precondition $ do + ensureWellTyped prog + prog <- runProgram prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/TrivialLetElimination.hs b/src/Jikka/Core/Convert/TrivialLetElimination.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/TrivialLetElimination.hs @@ -0,0 +1,93 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.TrivialLetElimination +-- Description : removes let-exprs whose variables are referenced at most only once. / その変数が高々 1 回しか参照されないような let 式を消去します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Convert.TrivialLetElimination + ( run, + run', + ) +where + +import Data.Functor +import qualified Data.Map as M +import Data.Maybe (fromMaybe) +import Jikka.Common.Error +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Lint + +plus :: Maybe Bool -> Maybe Bool -> Maybe Bool +plus (Just _) (Just _) = Just False +plus (Just p) Nothing = Just p +plus Nothing (Just p) = Just p +plus Nothing Nothing = Nothing + +isEliminatable :: VarName -> Expr -> Maybe Bool +isEliminatable x = \case + Var y -> if x == y then Just True else Nothing + Lit _ -> Nothing + App f e -> isEliminatable x f `plus` isEliminatable x e + Lam y _ e -> if x == y then Nothing else isEliminatable x e $> False -- moving an expr into a lambda may increase the time complexity + Let y _ e1 e2 -> isEliminatable x e1 `plus` (if x == y then Nothing else isEliminatable x e2) + +isEliminatableToplevelExpr :: VarName -> ToplevelExpr -> Maybe Bool +isEliminatableToplevelExpr x = \case + ResultExpr e -> isEliminatable x e + ToplevelLet y _ e cont -> isEliminatable x e `plus` (if x == y then Nothing else isEliminatableToplevelExpr x cont) + ToplevelLetRec f args _ body cont -> if x == f then Nothing else isEliminatableToplevelExpr x cont `plus` (if x `elem` map fst args then Nothing else isEliminatable x body) + +runExpr :: M.Map VarName Expr -> Expr -> Expr +runExpr env = \case + Var x -> fromMaybe (Var x) (M.lookup x env) + Lit lit -> Lit lit + App f e -> App (runExpr env f) (runExpr env e) + Lam x t body -> Lam x t (runExpr env body) + Let x t e1 e2 -> + let e1' = runExpr env e1 + in if isEliminatable x e2 /= Just False + then runExpr (M.insert x e1' env) e2 + else Let x t e1' (runExpr env e2) + +runToplevelExpr :: M.Map VarName Expr -> ToplevelExpr -> ToplevelExpr +runToplevelExpr env = \case + ResultExpr e -> ResultExpr (runExpr env e) + ToplevelLet x t e cont -> + let e' = runExpr env e + in if isEliminatableToplevelExpr x cont /= Just False + then runToplevelExpr (M.insert x e' env) cont + else ToplevelLet x t e' (runToplevelExpr env cont) + ToplevelLetRec f args ret body cont -> + ToplevelLetRec f args ret (runExpr env body) (runToplevelExpr env cont) + +run' :: Program -> Program +run' = runToplevelExpr M.empty + +-- | `run` remove let-exprs whose assigned variables are used only at most once. +-- This assumes that the program is alpha-converted. +-- +-- For example, this converts the following: +-- +-- > let f = fun y -> y +-- > in let x = 1 +-- > in f(x + x) +-- +-- to: +-- +-- > let x = 1 +-- > in (fun y -> y) (x + x) +-- +-- NOTE: this doesn't constant folding. +run :: MonadError Error m => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.ConstantPropagation" $ do + precondition $ do + ensureWellTyped prog + prog <- return $ run' prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Convert/TypeInfer.hs b/src/Jikka/Core/Convert/TypeInfer.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/TypeInfer.hs @@ -0,0 +1,230 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.TypeInfer +-- Description : does type inference. / 型推論を行います。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Convert.TypeInfer + ( run, + + -- * internal types and functions + Equation (..), + formularizeProgram, + sortEquations, + mergeAssertions, + Subst (..), + subst, + solveEquations, + substProgram, + ) +where + +import Control.Arrow (second) +import Control.Monad.State.Strict +import Control.Monad.Writer.Strict (MonadWriter, execWriterT, tell) +import qualified Data.Map.Strict as M +import Data.Monoid (Dual (..)) +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Format (formatType) +import Jikka.Core.Language.Expr +import Jikka.Core.Language.FreeVars +import Jikka.Core.Language.Lint +import Jikka.Core.Language.TypeCheck (literalToType, typecheckProgram) +import Jikka.Core.Language.Util + +data Equation + = TypeEquation Type Type + | TypeAssertion VarName Type + deriving (Eq, Ord, Show, Read) + +type Eqns = Dual [Equation] + +formularizeType :: MonadWriter Eqns m => Type -> Type -> m () +formularizeType t1 t2 = tell $ Dual [TypeEquation t1 t2] + +formularizeVarName :: MonadWriter Eqns m => VarName -> Type -> m () +formularizeVarName x t = tell $ Dual [TypeAssertion x t] + +formularizeExpr :: (MonadWriter Eqns m, MonadAlpha m) => Expr -> m Type +formularizeExpr = \case + Var x -> do + t <- genType + formularizeVarName x t + return t + Lit lit -> return $ literalToType lit + App f e -> do + ret <- genType + t <- formularizeExpr e + formularizeExpr' f (FunTy t ret) + return ret + Lam x t body -> do + formularizeVarName x t + ret <- formularizeExpr body + return $ FunTy t ret + Let x t e1 e2 -> do + formularizeVarName x t + formularizeExpr' e1 t + formularizeExpr e2 + +formularizeExpr' :: (MonadWriter Eqns m, MonadAlpha m) => Expr -> Type -> m () +formularizeExpr' e t = do + t' <- formularizeExpr e + formularizeType t t' + +formularizeToplevelExpr :: (MonadWriter Eqns m, MonadAlpha m) => ToplevelExpr -> m Type +formularizeToplevelExpr = \case + ResultExpr e -> formularizeExpr e + ToplevelLet x t e cont -> do + formularizeVarName x t + formularizeExpr' e t + formularizeToplevelExpr cont + ToplevelLetRec f args ret body cont -> do + formularizeVarName f (curryFunTy (map snd args) ret) + mapM_ (uncurry formularizeVarName) args + formularizeExpr' body ret + formularizeToplevelExpr cont + +formularizeProgram :: MonadAlpha m => Program -> m [Equation] +formularizeProgram prog = getDual <$> execWriterT (formularizeToplevelExpr prog) + +sortEquations :: [Equation] -> ([(Type, Type)], [(VarName, Type)]) +sortEquations = go [] [] + where + go eqns' assertions [] = (eqns', assertions) + go eqns' assertions (eqn : eqns) = case eqn of + TypeEquation t1 t2 -> go ((t1, t2) : eqns') assertions eqns + TypeAssertion x t -> go eqns' ((x, t) : assertions) eqns + +mergeAssertions :: [(VarName, Type)] -> [(Type, Type)] +mergeAssertions = go M.empty [] + where + go _ eqns [] = eqns + go gamma eqns ((x, t) : assertions) = case M.lookup x gamma of + Nothing -> go (M.insert x t gamma) eqns assertions + Just t' -> go gamma ((t, t') : eqns) assertions + +-- | `Subst` is type substituion. It's a mapping from type variables to their actual types. +newtype Subst = Subst {unSubst :: M.Map TypeName Type} + +subst :: Subst -> Type -> Type +subst sigma = \case + VarTy x -> + case M.lookup x (unSubst sigma) of + Nothing -> VarTy x + Just t -> subst sigma t + IntTy -> IntTy + BoolTy -> BoolTy + ListTy t -> ListTy (subst sigma t) + TupleTy ts -> TupleTy (map (subst sigma) ts) + FunTy t ret -> FunTy (subst sigma t) (subst sigma ret) + DataStructureTy ds -> DataStructureTy ds + +unifyTyVar :: (MonadState Subst m, MonadError Error m) => TypeName -> Type -> m () +unifyTyVar x t = + if x `elem` freeTyVars t + then throwInternalError $ "looped type equation " ++ unTypeName x ++ " = " ++ formatType t + else do + modify' (Subst . M.insert x t . unSubst) -- This doesn't introduce the loop. + +unifyType :: (MonadState Subst m, MonadError Error m) => Type -> Type -> m () +unifyType t1 t2 = wrapError' ("failed to unify " ++ formatType t1 ++ " and " ++ formatType t2) $ do + sigma <- get + t1 <- return $ subst sigma t1 -- shadowing + t2 <- return $ subst sigma t2 -- shadowing + case (t1, t2) of + _ | t1 == t2 -> return () + (VarTy x1, _) -> do + unifyTyVar x1 t2 + (_, VarTy x2) -> do + unifyTyVar x2 t1 + (ListTy t1, ListTy t2) -> do + unifyType t1 t2 + (TupleTy ts1, TupleTy ts2) -> do + if length ts1 == length ts2 + then mapM_ (uncurry unifyType) (zip ts1 ts2) + else throwInternalError $ "different type ctors " ++ formatType t1 ++ " and " ++ formatType t2 + (FunTy t1 ret1, FunTy t2 ret2) -> do + unifyType t1 t2 + unifyType ret1 ret2 + _ -> throwInternalError $ "different type ctors " ++ formatType t1 ++ " and " ++ formatType t2 + +solveEquations :: MonadError Error m => [(Type, Type)] -> m Subst +solveEquations eqns = wrapError' "failed to solve type equations" $ do + execStateT (mapM_ (uncurry unifyType) eqns) (Subst M.empty) + +-- | `substUnit` replaces all undetermined type variables with the unit type. +substUnit :: Type -> Type +substUnit = \case + VarTy _ -> TupleTy [] + IntTy -> IntTy + BoolTy -> BoolTy + ListTy t -> ListTy (substUnit t) + TupleTy ts -> TupleTy (map substUnit ts) + FunTy t ret -> FunTy (substUnit t) (substUnit ret) + DataStructureTy ds -> DataStructureTy ds + +-- | `subst'` does `subst` and replaces all undetermined type variables with the unit type. +subst' :: Subst -> Type -> Type +subst' sigma = substUnit . subst sigma + +substBuiltin :: Subst -> Builtin -> Builtin +substBuiltin sigma = mapTypeInBuiltin (subst' sigma) + +substLiteral :: Subst -> Literal -> Literal +substLiteral sigma = \case + LitBuiltin builtin -> LitBuiltin (substBuiltin sigma builtin) + LitInt n -> LitInt n + LitBool p -> LitBool p + LitNil t -> LitNil (subst' sigma t) + LitBottom t err -> LitBottom (subst' sigma t) err + +substExpr :: Subst -> Expr -> Expr +substExpr sigma = go + where + go = \case + Var x -> Var x + Lit lit -> Lit (substLiteral sigma lit) + App f e -> App (go f) (go e) + Lam x t body -> Lam x (subst' sigma t) (go body) + Let x t e1 e2 -> Let x (subst sigma t) (go e1) (go e2) + +substToplevelExpr :: Subst -> ToplevelExpr -> ToplevelExpr +substToplevelExpr sigma = \case + ResultExpr e -> ResultExpr (substExpr sigma e) + ToplevelLet x t e cont -> ToplevelLet x (subst' sigma t) (substExpr sigma e) (substToplevelExpr sigma cont) + ToplevelLetRec f args ret body cont -> ToplevelLetRec f (map (second (subst' sigma)) args) (subst' sigma ret) (substExpr sigma body) (substToplevelExpr sigma cont) + +substProgram :: Subst -> Program -> Program +substProgram = substToplevelExpr + +-- | `run` does type inference. +-- +-- * This assumes that program has no name conflicts. +-- +-- Before: +-- +-- > let f = fun y -> y +-- > in let x = 1 +-- > in f(x + x) +-- +-- After: +-- +-- > let f: int -> int = fun y: int -> y +-- > in let x: int = 1 +-- > in f(x + x) +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.TypeInfer" $ do + eqns <- formularizeProgram prog + let (eqns', assertions) = sortEquations eqns + let eqns'' = mergeAssertions assertions + sigma <- solveEquations (eqns' ++ eqns'') + prog <- return $ substProgram sigma prog + postcondition $ do + typecheckProgram prog + return prog diff --git a/src/Jikka/Core/Convert/UnpackTuple.hs b/src/Jikka/Core/Convert/UnpackTuple.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Convert/UnpackTuple.hs @@ -0,0 +1,91 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Convert.UnpackTuples +-- Description : unpacks and flattens tuples. / タプルを展開し平坦にします。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Convert.UnpackTuple + ( run, + + -- * internal rules + rule, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import qualified Jikka.Core.Convert.Alpha as Alpha +import Jikka.Core.Language.Beta +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Lint +import Jikka.Core.Language.RewriteRules +import Jikka.Core.Language.Util + +rule :: (MonadAlpha m, MonadError Error m) => RewriteRule m +rule = + let return' = return . Just + in RewriteRule $ \_ -> \case + App (Lam x (TupleTy ts) body) e -> case curryApp e of + (Tuple' ts', es) -> do + when (ts /= ts') $ do + throwInternalError "the types of tuple don't match" + when (length ts /= length es) $ do + throwInternalError "the sizes of tuple don't match" + xs <- replicateM (length ts) (genVarName x) + body' <- substitute x (uncurryApp (Tuple' ts) (map Var xs)) body + return' $ uncurryApp (curryLam (zip xs ts) body') es + _ -> return Nothing + App (Tuple' [_]) (Proj' [_] 0 e) -> return' e + Proj' ts i e -> case curryApp e of + (Tuple' _, es) -> return' $ es !! i + (Lit (LitBuiltin (If _)), [e1, e2, e3]) -> return' $ If' (ts !! i) e1 (Proj' ts i e2) (Proj' ts i e3) + _ -> return Nothing + Foldl' t2 (TupleTy [t1]) (Lam x1 (TupleTy [_]) (Lam x2 _ body)) e es -> do + body' <- substitute x1 (App (Tuple' [t1]) (Var x1)) (Proj' [t1] 0 body) + return' $ App (Tuple' [t1]) (Foldl' t2 t1 (Lam2 x1 t1 x2 t2 body') (Proj' [t1] 0 e) es) + Scanl' t2 (TupleTy [t1]) (Lam x1 _ (Lam x2 (TupleTy [_]) body)) e es -> do + body' <- substitute x1 (App (Tuple' [t1]) (Var x1)) (Proj' [t1] 0 body) + let e' = Scanl' t2 t1 (Lam2 x1 t1 x2 t2 body') (Proj' [t1] 0 e) es + y <- genVarName' + let f = Map' t1 (TupleTy [t1]) (Lam y t1 (App (Tuple' [t1]) (Var y))) + return' $ f e' + Iterate' (TupleTy [t]) n (Lam x (TupleTy [_]) body) base -> do + body' <- substitute x (App (Tuple' [t]) (Var x)) (Proj' [t] 0 body) + return' $ uncurryApp (Tuple' [t]) [Iterate' t n (Lam x t body') (Proj' [t] 0 base)] + _ -> return Nothing + +runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram = applyRewriteRuleProgram' rule + +-- | `run` removes unnecessary introductions and eliminations of tuples. +-- For example, this converts the following: +-- +-- > (fun xs -> (proj0 xs) + (proj1 xs)) (tuple 2 1) +-- +-- to the follwoing: +-- +-- > (fun x0 x1 -> x0 + x1) 2 1 +-- +-- This can remove 1-tuples over higher-order functions. +-- For example, this converts the following: +-- +-- > foldl (fun xs y -> tuple (proj0 xs + y) (tuple 0) [1, 2, 3] +-- +-- to the follwoing: +-- +-- > tuple (foldl (fun x y -> x + y) 0 [1, 2, 3]) +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.Core.Convert.UnpackTuple" $ do + precondition $ do + ensureWellTyped prog + prog <- Alpha.run prog + prog <- runProgram prog + postcondition $ do + ensureWellTyped prog + return prog diff --git a/src/Jikka/Core/Evaluate.hs b/src/Jikka/Core/Evaluate.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Evaluate.hs @@ -0,0 +1,320 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TupleSections #-} + +-- | +-- Module : Jikka.Core.Evaluate +-- Description : executes the expr of our core language. / core 言語の式を実行します。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- `Jikka.Core.Evaluate` evaluates exprs to values. Also this recognizes users' inputs at once. +-- +-- The implementation assumes that all variable names don't conflict even when their scopes are distinct. +module Jikka.Core.Evaluate + ( run, + callProgram, + Value (..), + ) +where + +import Control.Monad.Except +import Data.Bits +import Data.List (maximumBy, minimumBy, sortBy) +import qualified Data.Vector as V +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Common.Matrix +import Jikka.Core.Format (formatBuiltinIsolated) +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Lint +import Jikka.Core.Language.Runtime +import Jikka.Core.Language.Util +import Jikka.Core.Language.Value + +-- ----------------------------------------------------------------------------- +-- builtins + +iterate' :: MonadError Error m => Integer -> Value -> Value -> m Value +iterate' n _ _ | n < 0 = throwRuntimeError $ "negative number of iteration: " ++ show n +iterate' 0 _ base = return base +iterate' n step base = do + base <- callValue step [base] + iterate' (n - 1) step base + +map' :: MonadError Error m => Value -> V.Vector Value -> m (V.Vector Value) +map' f a = V.fromList <$> mapM (\val -> callValue f [val]) (V.toList a) + +scanM :: Monad m => (a -> b -> m a) -> a -> V.Vector b -> m (V.Vector a) +scanM f y xs = do + (ys, y) <- V.foldM (\(ys, y) x -> (y : ys,) <$> f y x) ([], y) xs + return $ V.fromList (reverse (y : ys)) + +atEither :: MonadError Error m => V.Vector a -> Integer -> m a +atEither xs i = case xs V.!? fromInteger i of + Just x -> return x + Nothing -> throwRuntimeError $ "out of bounds: length = " ++ show (V.length xs) ++ ", index = " ++ show i + +setAtEither :: MonadError Error m => V.Vector a -> Integer -> a -> m (V.Vector a) +setAtEither xs i x = + if 0 <= i && i < fromIntegral (V.length xs) + then return $ xs V.// [(fromInteger i, x)] + else throwRuntimeError $ "out of bounds: length = " ++ show (V.length xs) ++ ", index = " ++ show i + +sortVector :: V.Vector Value -> V.Vector Value +sortVector = V.fromList . sortBy compareValues' . V.toList + +range1 :: MonadError Error m => Integer -> m (V.Vector Value) +range1 n | n < 0 = throwRuntimeError $ "invalid argument for range1: " ++ show n +range1 n = return $ V.fromList (map ValInt [0 .. n - 1]) + +range2 :: MonadError Error m => Integer -> Integer -> m (V.Vector Value) +range2 l r | l > r = throwRuntimeError $ "invalid argument for range2: " ++ show (l, r) +range2 l r = return $ V.fromList (map ValInt [l .. r - 1]) + +range3 :: MonadError Error m => Integer -> Integer -> Integer -> m (V.Vector Value) +range3 l r step | not (l <= r && step >= 0) = throwRuntimeError $ "invalid argument for range3: " ++ show (l, r, step) +range3 l r step = return $ V.fromList (map ValInt [l, l + step .. r]) + +matap' :: (Num a, MonadError Error m) => Matrix a -> V.Vector a -> m (V.Vector a) +matap' f x | snd (matsize f) /= V.length x = throwInternalError "invalid argument" +matap' f x = return $ matap f x + +matadd' :: (Num a, MonadError Error m) => Matrix a -> Matrix a -> m (Matrix a) +matadd' f g | matsize f /= matsize g = throwInternalError "invalid argument" +matadd' f g = return $ matadd f g + +matmul' :: (Num a, MonadError Error m) => Matrix a -> Matrix a -> m (Matrix a) +matmul' f g | snd (matsize f) /= fst (matsize g) = throwInternalError "invalid argument" +matmul' f g = return $ matmul f g + +matpow' :: (Num a, Show a, MonadError Error m) => Matrix a -> Integer -> m (Matrix a) +matpow' f _ | let (h, w) = matsize f in h /= w = throwInternalError $ "matrix is not square: " ++ show (matsize f) +matpow' _ k | k < 0 = throwRuntimeError $ "exponent is negative: " ++ show k +matpow' f k = return $ matpow f k + +convexHullTrickGetMin :: MonadError Error m => V.Vector (Integer, Integer) -> Integer -> m Integer +convexHullTrickGetMin cht x = + if V.null cht + then throwRuntimeError "the set of lines is empty" + else return $ V.minimum (V.map (\(a, b) -> a * x + b) cht) + +segmentTreeGetRange :: MonadError Error m => Semigroup' -> [Integer] -> Integer -> Integer -> m Integer +segmentTreeGetRange semigrp segtree l r + | l > r = throwRuntimeError $ "the range has negative length: l = " ++ show l ++ ", r = " ++ show r + | l == r = throwRuntimeError $ "the range is empty: l = r = " ++ show l + | otherwise = + let slice = take (fromInteger (r - l)) (drop (fromInteger l) segtree) + in return $ case semigrp of + SemigroupIntPlus -> sum slice + SemigroupIntMin -> minimum slice + SemigroupIntMax -> maximum slice + +build :: MonadError Error m => (V.Vector Value -> m Value) -> V.Vector Value -> Integer -> m (V.Vector Value) +build _ _ n | n < 0 = throwRuntimeError $ "negative length: " ++ show n +build _ xs 0 = return xs +build f xs n = do + y <- f xs + build f (V.snoc xs y) (n - 1) + +-- ----------------------------------------------------------------------------- +-- evaluator + +callBuiltin :: MonadError Error m => Builtin -> [Value] -> m Value +callBuiltin builtin args = wrapError' ("while calling builtin " ++ formatBuiltinIsolated builtin) $ do + let go0 ret f = callValue (ret f) args + let go1' t1 ret f = case args of + v1 : args -> do + f <- ret <$> (f =<< t1 v1) + callValue f args + _ -> return $ ValBuiltin builtin args + let go1 t1 ret f = go1' t1 ret (return . f) + let go2' t1 t2 ret f = case args of + v1 : v2 : args -> do + f <- ret <$> join (f <$> t1 v1 <*> t2 v2) + callValue f args + _ -> return $ ValBuiltin builtin args + let go2 t1 t2 ret f = go2' t1 t2 ret ((return .) . f) + let go3' t1 t2 t3 ret f = case args of + v1 : v2 : v3 : args -> do + f <- ret <$> join (f <$> t1 v1 <*> t2 v2 <*> t3 v3) + callValue f args + _ -> return $ ValBuiltin builtin args + let go3 t1 t2 t3 ret f = go3' t1 t2 t3 ret (((return .) .) . f) + let goN n t ret f = + if length args < n + then return $ ValBuiltin builtin args + else do + f <- ret . f <$> mapM t (take n args) + callValue f (drop n args) + case builtin of + -- arithmetical functions + Negate -> go1 valueToInt ValInt negate + Plus -> go2 valueToInt valueToInt ValInt (+) + Minus -> go2 valueToInt valueToInt ValInt (-) + Mult -> go2 valueToInt valueToInt ValInt (*) + FloorDiv -> go2' valueToInt valueToInt ValInt floorDiv + FloorMod -> go2' valueToInt valueToInt ValInt floorMod + CeilDiv -> go2' valueToInt valueToInt ValInt ceilDiv + CeilMod -> go2' valueToInt valueToInt ValInt ceilMod + Pow -> go2 valueToInt valueToInt ValInt (^) + -- advanced arithmetical functions + Abs -> go1 valueToInt ValInt abs + Gcd -> go2 valueToInt valueToInt ValInt gcd + Lcm -> go2 valueToInt valueToInt ValInt lcm + Min2 _ -> go2 pure pure id minValue + Max2 _ -> go2 pure pure id maxValue + Iterate _ -> go3' valueToInt pure pure id $ \n step base -> iterate' n step base + -- logical functions + Not -> go1 valueToBool ValBool not + And -> go2 valueToBool valueToBool ValBool (&&) + Or -> go2 valueToBool valueToBool ValBool (||) + Implies -> go2 valueToBool valueToBool ValBool $ \p q -> not p || q + If _ -> go3 valueToBool pure pure id $ \p a b -> if p then a else b + -- bitwise functions + BitNot -> go1 valueToInt ValInt complement + BitAnd -> go2 valueToInt valueToInt ValInt (.&.) + BitOr -> go2 valueToInt valueToInt ValInt (.|.) + BitXor -> go2 valueToInt valueToInt ValInt xor + BitLeftShift -> go2 valueToInt valueToInt ValInt $ \a b -> a `shift` fromInteger b + BitRightShift -> go2 valueToInt valueToInt ValInt $ \a b -> a `shift` fromInteger (- b) + -- matrix functions + MatAp _ _ -> go2' valueToMatrix valueToVector valueFromVector matap' + MatZero n -> go0 valueFromMatrix (matzero n) + MatOne n -> go0 valueFromMatrix (matone n) + MatAdd _ _ -> go2' valueToMatrix valueToMatrix valueFromMatrix matadd' + MatMul _ _ _ -> go2' valueToMatrix valueToMatrix valueFromMatrix matmul' + MatPow _ -> go2' valueToMatrix valueToInt valueFromMatrix matpow' + VecFloorMod _ -> go2 valueToVector valueToInt valueFromVector $ \x m -> V.map (`mod` m) x + MatFloorMod _ _ -> go2 valueToMatrix valueToInt valueFromMatrix $ \f m -> fmap (`mod` m) f + -- modular functions + ModNegate -> go2 valueToInt valueToInt ValInt $ \a m -> (- a) `mod` m + ModPlus -> go3 valueToInt valueToInt valueToInt ValInt $ \a b m -> (a + b) `mod` m + ModMinus -> go3 valueToInt valueToInt valueToInt ValInt $ \a b m -> (a - b) `mod` m + ModMult -> go3 valueToInt valueToInt valueToInt ValInt $ \a b m -> (a * b) `mod` m + ModInv -> go2' valueToInt valueToInt ValInt modinv + ModPow -> go3' valueToInt valueToInt valueToInt ValInt modpow + ModMatAp _ _ -> go3' pure pure valueToInt valueFromModVector $ \f x m -> join (matap' <$> valueToModMatrix m f <*> valueToModVector m x) + ModMatAdd _ _ -> go3' pure pure valueToInt valueFromModMatrix $ \f g m -> join (matadd' <$> valueToModMatrix m f <*> valueToModMatrix m g) + ModMatMul _ _ _ -> go3' pure pure valueToInt valueFromModMatrix $ \f g m -> join (matmul' <$> valueToModMatrix m f <*> valueToModMatrix m g) + ModMatPow _ -> go3' pure valueToInt valueToInt valueFromModMatrix $ \f k m -> join (matpow' <$> valueToModMatrix m f <*> pure k) + -- list functions + Cons _ -> go2 pure valueToList ValList V.cons + Snoc _ -> go2 valueToList pure ValList V.snoc + Foldl _ _ -> go3' pure pure valueToList id $ \f x a -> V.foldM (\x y -> callValue f [x, y]) x a + Scanl _ _ -> go3' pure pure valueToList ValList $ \f x a -> scanM (\x y -> callValue f [x, y]) x a + Build _ -> go3' pure valueToList valueToInt ValList $ \f xs n -> build (\xs -> callValue f [ValList xs]) xs n + Len _ -> go1 valueToList ValInt (fromIntegral . V.length) + Map _ _ -> go2' pure valueToList ValList map' + Filter _ -> go2' pure valueToList ValList $ \f xs -> V.filterM (\x -> (/= ValBool False) <$> callValue f [x]) xs + At _ -> go2' valueToList valueToInt id atEither + SetAt _ -> go3' valueToList valueToInt pure ValList setAtEither + Elem _ -> go2 pure valueToList ValBool V.elem + Sum -> go1 valueToIntList ValInt sum + ModSum -> go2 valueToIntList valueToInt ValInt $ \xs m -> sum xs `mod` m + Product -> go1 valueToIntList ValInt product + ModProduct -> go2 valueToIntList valueToInt ValInt $ \xs m -> product xs `mod` m + Min1 _ -> go1 valueToList id (V.minimumBy compareValues') + Max1 _ -> go1 valueToList id (V.maximumBy compareValues') + ArgMin _ -> go1 valueToList ValInt $ \xs -> snd (minimumBy (\(x, i) (y, j) -> compareValues' x y <> compare i j) (zip (V.toList xs) [0 ..])) + ArgMax _ -> go1 valueToList ValInt $ \xs -> snd (maximumBy (\(x, i) (y, j) -> compareValues' x y <> compare i j) (zip (V.toList xs) [0 ..])) + All -> go1 valueToBoolList ValBool and + Any -> go1 valueToBoolList ValBool or + Sorted _ -> go1 valueToList ValList sortVector + Reversed _ -> go1 valueToList ValList V.reverse + Range1 -> go1' valueToInt ValList range1 + Range2 -> go2' valueToInt valueToInt ValList range2 + Range3 -> go3' valueToInt valueToInt valueToInt ValList range3 + -- tuple functions + Tuple ts -> goN (length ts) pure ValTuple id + Proj _ n -> go1 valueToTuple id (!! n) + -- -- comparison + LessThan _ -> go2 pure pure ValBool $ \a b -> compareValues a b == Just LT + LessEqual _ -> go2 pure pure ValBool $ \a b -> compareValues a b /= Just GT + GreaterThan _ -> go2 pure pure ValBool $ \a b -> compareValues a b == Just GT + GreaterEqual _ -> go2 pure pure ValBool $ \a b -> compareValues a b /= Just LT + Equal _ -> go2 pure pure ValBool (==) + NotEqual _ -> go2 pure pure ValBool (/=) + -- combinational functions + Fact -> go1' valueToInt ValInt fact + Choose -> go2' valueToInt valueToInt ValInt choose + Permute -> go2' valueToInt valueToInt ValInt permute + MultiChoose -> go2' valueToInt valueToInt ValInt multichoose + -- data structures + ConvexHullTrickInit -> go0 ValList V.empty + ConvexHullTrickGetMin -> go2' (V.mapM valueToIntPair <=< valueToList) valueToInt ValInt convexHullTrickGetMin + ConvexHullTrickInsert -> go3 valueToList pure pure ValList $ \cht a b -> V.snoc cht (ValTuple [a, b]) + SegmentTreeInitList _ -> go1 pure id id + SegmentTreeGetRange semigrp -> go3' valueToIntList valueToInt valueToInt ValInt (segmentTreeGetRange semigrp) + SegmentTreeSetPoint _ -> go3' valueToList valueToInt pure ValList setAtEither + +callLambda :: MonadError Error m => Maybe VarName -> Env -> VarName -> Type -> Expr -> [Value] -> m Value +callLambda = \name env x t body args -> wrapError' ("while calling lambda " ++ maybe "(anonymous)" unVarName name) $ go Nothing env x t body args + where + go name env x t body [] = return $ ValLambda name env x t body + go name env x _ body (e : args) = maybe id (\name -> wrapError' $ "while calling lambda " ++ unVarName name) name $ do + body <- evaluateExpr ((x, e) : env) body + case body of + ValLambda name env x t body -> go name env x t body args + _ -> callValue body args + +callValue :: MonadError Error m => Value -> [Value] -> m Value +callValue f args = case (f, args) of + (ValBuiltin builtin args', _) -> callBuiltin builtin (args' ++ args) + (ValLambda name env x t body, _) -> callLambda name env x t body args + (_, []) -> return f + _ -> throwInternalError $ "cannot call a non-function: " ++ formatValue f + +evaluateExpr :: MonadError Error m => Env -> Expr -> m Value +evaluateExpr env = \case + Var x -> case lookup x env of + Nothing -> throwInternalError $ "undefined variable: " ++ unVarName x + Just val -> return val + Lit lit -> literalToValue lit + If' _ p e1 e2 -> do + p <- valueToBool =<< evaluateExpr env p + if p + then evaluateExpr env e1 + else evaluateExpr env e2 + e@(App _ _) -> do + let (f, args) = curryApp e + f <- evaluateExpr env f + args <- mapM (evaluateExpr env) args + callValue f args + Lam x t body -> return $ ValLambda Nothing env x t body + Let x _ e1 e2 -> do + v1 <- evaluateExpr env e1 + evaluateExpr ((x, v1) : env) e2 + +callToplevelExpr :: (MonadFix m, MonadError Error m) => Env -> ToplevelExpr -> [Value] -> m Value +callToplevelExpr env e args = case e of + ToplevelLet x _ e cont -> do + val <- evaluateExpr env e + callToplevelExpr ((x, val) : env) cont args + ToplevelLetRec f args' _ body cont -> do + val <- mfix $ \val -> evaluateExpr ((f, val) : env) (curryLam args' body) + callToplevelExpr ((f, val) : env) cont args + ResultExpr e -> do + val <- evaluateExpr env e + callValue val args + +-- | `callProgram` evaluates programs with given arguments. +-- This function assumes that given programs are ready for eager evaluation (@ensureEagerlyEvaluatable@). +callProgram :: (MonadFix m, MonadError Error m) => Program -> [Value] -> m Value +callProgram prog args = wrapError' "Jikka.Core.Evaluate" $ do + precondition $ do + ensureEagerlyEvaluatable prog + ensureWellTyped prog + callToplevelExpr [] prog args + +-- ----------------------------------------------------------------------------- +-- run + +run :: (MonadAlpha m, MonadFix m, MonadError Error m) => Program -> [Value] -> m Value +run prog args = do + callProgram prog args diff --git a/src/Jikka/Core/Format.hs b/src/Jikka/Core/Format.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Format.hs @@ -0,0 +1,252 @@ +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Format +-- Description : converts the syntax trees of core language to strings. / core 言語の構文木を文字列に変換します。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- TODO: add parens with considering precedences. +module Jikka.Core.Format + ( run, + formatBuiltinIsolated, + formatBuiltin, + formatType, + formatExpr, + formatProgram, + ) +where + +import Data.Char (toLower) +import Data.List (intercalate) +import Data.Text (Text, pack) +import Jikka.Common.Format.AutoIndent +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Util + +paren :: String -> String +paren s = "(" ++ s ++ ")" + +formatType :: Type -> String +formatType = \case + VarTy (TypeName a) -> a + IntTy -> "int" + BoolTy -> "bool" + ListTy t -> formatType t ++ " list" + TupleTy ts -> case ts of + [t] -> paren $ formatType t ++ "," + _ -> paren $ intercalate " * " (map formatType ts) + t@(FunTy _ _) -> + let (ts, ret) = uncurryFunTy t + in paren $ intercalate " -> " (map formatType (ts ++ [ret])) + DataStructureTy ds -> formatDataStructure ds + +formatDataStructure :: DataStructure -> String +formatDataStructure = \case + ConvexHullTrick -> "convex-hull-trick" + SegmentTree semigrp -> "segment-tree<" ++ formatSemigroup semigrp ++ ">" + +formatSemigroup :: Semigroup' -> String +formatSemigroup = \case + SemigroupIntPlus -> "int.plus" + SemigroupIntMin -> "int.min" + SemigroupIntMax -> "int.max" + +data Builtin' + = Fun [Type] String + | PrefixOp String + | InfixOp [Type] String + | At' Type + | If' Type + deriving (Eq, Ord, Show, Read) + +fun :: String -> Builtin' +fun = Fun [] + +infixOp :: String -> Builtin' +infixOp = InfixOp [] + +analyzeBuiltin :: Builtin -> Builtin' +analyzeBuiltin = \case + -- arithmetical functions + Negate -> PrefixOp "negate" + Plus -> infixOp "+" + Minus -> infixOp "-" + Mult -> infixOp "*" + FloorDiv -> infixOp "/" + FloorMod -> infixOp "%" + CeilDiv -> fun "ceildiv" + CeilMod -> fun "ceilmod" + Pow -> infixOp "**" + -- advanced arithmetical functions + Abs -> fun "abs" + Gcd -> fun "gcd" + Lcm -> fun "lcm" + Min2 t -> Fun [t] "min" + Max2 t -> Fun [t] "max" + -- logical functions + Not -> PrefixOp "not" + And -> infixOp "and" + Or -> infixOp "or" + Implies -> infixOp "implies" + If t -> If' t + -- bitwise functions + BitNot -> PrefixOp "~" + BitAnd -> infixOp "&" + BitOr -> infixOp "|" + BitXor -> infixOp "^" + BitLeftShift -> infixOp "<<" + BitRightShift -> infixOp ">>" + -- matrix functions + MatAp _ _ -> fun "matap" + MatZero _ -> fun "matzero" + MatOne _ -> fun "matone" + MatAdd _ _ -> fun "matadd" + MatMul _ _ _ -> fun "matmul" + MatPow _ -> fun "matpow" + VecFloorMod _ -> fun "vecfloormod" + MatFloorMod _ _ -> fun "matfloormod" + -- modular functions + ModNegate -> fun "modnegate" + ModPlus -> fun "modplus" + ModMinus -> fun "modminus" + ModMult -> fun "modmult" + ModInv -> fun "modinv" + ModPow -> fun "modpow" + ModMatAp _ _ -> fun "modmatap" + ModMatAdd _ _ -> fun "modmatadd" + ModMatMul _ _ _ -> fun "modmatmul" + ModMatPow _ -> fun "modmatpow" + -- list functions + Cons t -> Fun [t] "cons" + Snoc t -> Fun [t] "snoc" + Foldl t1 t2 -> Fun [t1, t2] "foldl" + Scanl t1 t2 -> Fun [t1, t2] "scanl" + Build t -> Fun [t] "build" + Iterate t -> Fun [t] "iterate" + Len t -> Fun [t] "len" + Map t1 t2 -> Fun [t1, t2] "map" + Filter t -> Fun [t] "filter" + At t -> At' t + SetAt t -> Fun [t] "setAt" + Elem t -> Fun [t] "elem" + Sum -> fun "sum" + Product -> fun "product" + ModSum -> fun "modsum" + ModProduct -> fun "modproduct" + Min1 t -> Fun [t] "min1" + Max1 t -> Fun [t] "max1" + ArgMin t -> Fun [t] "argmin" + ArgMax t -> Fun [t] "argmax" + All -> fun "all" + Any -> fun "any" + Sorted t -> Fun [t] "sort" + Reversed t -> Fun [t] "reverse" + Range1 -> fun "range1" + Range2 -> fun "range2" + Range3 -> fun "range3" + -- tuple functions + Tuple ts -> Fun ts "tuple" + Proj ts n -> Fun ts ("proj" ++ show n) + -- comparison + LessThan t -> InfixOp [t] "<" + LessEqual t -> InfixOp [t] "<=" + GreaterThan t -> InfixOp [t] ">" + GreaterEqual t -> InfixOp [t] ">=" + Equal t -> InfixOp [t] "==" + NotEqual t -> InfixOp [t] "!=" + -- combinational functions + Fact -> fun "fact" + Choose -> fun "choose" + Permute -> fun "permute" + MultiChoose -> fun "multichoose" + -- data structures + ConvexHullTrickInit -> fun "cht.init" + ConvexHullTrickGetMin -> fun "cht.getmin" + ConvexHullTrickInsert -> fun "cht.insert" + SegmentTreeInitList _ -> fun "segtree.initlist" + SegmentTreeGetRange _ -> fun "segtree.getrange" + SegmentTreeSetPoint _ -> fun "segtree.setpoint" + +formatTemplate :: [Type] -> String +formatTemplate = \case + [] -> "" + ts -> "<" ++ intercalate ", " (map formatType ts) ++ ">" + +formatFunCall :: String -> [Expr] -> String +formatFunCall f = \case + [] -> f + args -> f ++ "(" ++ intercalate ", " (map formatExpr' args) ++ ")" + +formatBuiltinIsolated' :: Builtin' -> String +formatBuiltinIsolated' = \case + Fun ts name -> name ++ formatTemplate ts + PrefixOp op -> paren op + InfixOp ts op -> paren $ op ++ formatTemplate ts + At' t -> paren $ "at" ++ formatTemplate [t] + If' t -> paren $ "if-then-else" ++ formatTemplate [t] + +formatBuiltinIsolated :: Builtin -> String +formatBuiltinIsolated = formatBuiltinIsolated' . analyzeBuiltin + +formatBuiltin' :: Builtin' -> [Expr] -> String +formatBuiltin' builtin args = case (builtin, args) of + (Fun _ name, _) -> formatFunCall name args + (PrefixOp op, e1 : args) -> formatFunCall (paren $ op ++ " " ++ formatExpr' e1) args + (InfixOp _ op, e1 : e2 : args) -> formatFunCall (paren $ formatExpr' e1 ++ " " ++ op ++ " " ++ formatExpr' e2) args + (At' _, e1 : e2 : args) -> formatFunCall (paren $ formatExpr' e1 ++ ")[" ++ formatExpr' e2 ++ "]") args + (If' _, e1 : e2 : e3 : args) -> formatFunCall (paren $ "if" ++ " " ++ formatExpr' e1 ++ " then " ++ formatExpr' e2 ++ " else " ++ formatExpr' e3) args + _ -> formatFunCall (formatBuiltinIsolated' builtin) args + +formatBuiltin :: Builtin -> [Expr] -> String +formatBuiltin = formatBuiltin' . analyzeBuiltin + +formatLiteral :: Literal -> String +formatLiteral = \case + LitBuiltin builtin -> formatBuiltinIsolated builtin + LitInt n -> show n + LitBool p -> map toLower $ show p + LitNil t -> "nil" ++ formatTemplate [t] + LitBottom t _ -> "bottom" ++ formatTemplate [t] + +formatFormalArgs :: [(VarName, Type)] -> String +formatFormalArgs args = unwords $ map (\(x, t) -> paren (unVarName x ++ ": " ++ formatType t)) args + +formatExpr' :: Expr -> String +formatExpr' = \case + Var x -> unVarName x + Lit lit -> formatLiteral lit + e@(App _ _) -> + let (f, args) = curryApp e + in case f of + Var x -> formatFunCall (unVarName x) args + Lit (LitBuiltin builtin) -> formatBuiltin builtin args + _ -> formatFunCall (formatExpr' f) args + e@(Lam _ _ _) -> + let (args, body) = uncurryLam e + in paren $ "fun " ++ formatFormalArgs args ++ " ->\n" ++ indent ++ "\n" ++ formatExpr' body ++ "\n" ++ dedent ++ "\n" + Let x t e1 e2 -> "let " ++ unVarName x ++ ": " ++ formatType t ++ " =\n" ++ indent ++ "\n" ++ formatExpr' e1 ++ "\n" ++ dedent ++ "\nin " ++ formatExpr' e2 + +formatExpr :: Expr -> String +formatExpr = unwords . makeIndentFromMarkers 4 . lines . formatExpr' + +formatToplevelExpr :: ToplevelExpr -> [String] +formatToplevelExpr = \case + ResultExpr e -> lines (formatExpr' e) + ToplevelLet x t e cont -> let' (unVarName x) t e cont + ToplevelLetRec f args ret e cont -> let' ("rec " ++ unVarName f ++ " " ++ formatFormalArgs args) ret e cont + where + let' s t e cont = + ["let " ++ s ++ ": " ++ formatType t ++ " =", indent] + ++ lines (formatExpr' e) + ++ [dedent, "in"] + ++ formatToplevelExpr cont + +formatProgram :: Program -> String +formatProgram = unlines . makeIndentFromMarkers 4 . formatToplevelExpr + +run :: Applicative m => Program -> m Text +run = pure . pack . formatProgram diff --git a/src/Jikka/Core/Language/ArithmeticalExpr.hs b/src/Jikka/Core/Language/ArithmeticalExpr.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Language/ArithmeticalExpr.hs @@ -0,0 +1,277 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE TupleSections #-} + +module Jikka.Core.Language.ArithmeticalExpr + ( -- * Basic functions + ArithmeticalExpr, + parseArithmeticalExpr, + formatArithmeticalExpr, + integerArithmeticalExpr, + negateArithmeticalExpr, + plusArithmeticalExpr, + minusArithmeticalExpr, + multArithmeticalExpr, + isZeroArithmeticalExpr, + isOneArithmeticalExpr, + + -- * Advanced functions + unNPlusKPattern, + makeVectorFromArithmeticalExpr, + makeAffineFunctionFromArithmeticalExpr, + splitConstantFactorArithmeticalExpr, + ) +where + +import Control.Arrow +import Control.Monad +import Control.Monad.ST +import Control.Monad.Trans +import Control.Monad.Trans.Maybe +import Data.List (findIndices, groupBy, sort, sortBy) +import Data.STRef +import qualified Data.Vector as V +import qualified Data.Vector.Mutable as MV +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.FreeVars + +data ProductExpr = ProductExpr + { productExprConst :: Integer, + productExprList :: [Expr] + } + deriving (Eq, Ord, Show, Read) + +data SumExpr = SumExpr + { sumExprList :: [ProductExpr], + sumExprConst :: Integer + } + deriving (Eq, Ord, Show, Read) + +newtype ArithmeticalExpr = ArithmeticalExpr {unArithmeticalExpr :: SumExpr} + deriving (Show) + +instance Eq ArithmeticalExpr where + e1 == e2 = unArithmeticalExpr (normalizeArithmeticalExpr e1) == unArithmeticalExpr (normalizeArithmeticalExpr e2) + +instance Ord ArithmeticalExpr where + e1 `compare` e2 = unArithmeticalExpr (normalizeArithmeticalExpr e1) `compare` unArithmeticalExpr (normalizeArithmeticalExpr e2) + +integerProductExpr :: Integer -> ProductExpr +integerProductExpr n = + ProductExpr + { productExprConst = n, + productExprList = [] + } + +negateProductExpr :: ProductExpr -> ProductExpr +negateProductExpr e = e {productExprConst = negate (productExprConst e)} + +multProductExpr :: ProductExpr -> ProductExpr -> ProductExpr +multProductExpr e1 e2 = + ProductExpr + { productExprConst = productExprConst e1 * productExprConst e2, + productExprList = productExprList e1 ++ productExprList e2 + } + +parseProductExpr :: Expr -> ProductExpr +parseProductExpr = \case + LitInt' n -> ProductExpr {productExprConst = n, productExprList = []} + Negate' e -> negateProductExpr (parseProductExpr e) + Mult' e1 e2 -> multProductExpr (parseProductExpr e1) (parseProductExpr e2) + Pow' e1 (LitInt' k) | 0 <= k && k < 10 -> iterate (multProductExpr (parseProductExpr e1)) (integerProductExpr 1) !! fromInteger k + e -> ProductExpr {productExprConst = 1, productExprList = [e]} + +sumExprFromProductExpr :: ProductExpr -> SumExpr +sumExprFromProductExpr e = + SumExpr + { sumExprList = [e], + sumExprConst = 0 + } + +arithmeticalExprFromProductExpr :: ProductExpr -> ArithmeticalExpr +arithmeticalExprFromProductExpr = ArithmeticalExpr . sumExprFromProductExpr + +integerSumExpr :: Integer -> SumExpr +integerSumExpr n = + SumExpr + { sumExprConst = n, + sumExprList = [] + } + +integerArithmeticalExpr :: Integer -> ArithmeticalExpr +integerArithmeticalExpr = ArithmeticalExpr . integerSumExpr + +negateSumExpr :: SumExpr -> SumExpr +negateSumExpr e = + SumExpr + { sumExprList = map negateProductExpr (sumExprList e), + sumExprConst = negate (sumExprConst e) + } + +plusSumExpr :: SumExpr -> SumExpr -> SumExpr +plusSumExpr e1 e2 = + SumExpr + { sumExprList = sumExprList e1 ++ sumExprList e2, + sumExprConst = sumExprConst e1 + sumExprConst e2 + } + +multSumExpr :: SumExpr -> SumExpr -> SumExpr +multSumExpr e1 e2 = + SumExpr + { sumExprList = + let es1 = parseProductExpr (LitInt' (sumExprConst e1)) : sumExprList e1 + es2 = parseProductExpr (LitInt' (sumExprConst e2)) : sumExprList e2 + in map (uncurry multProductExpr) ((,) <$> es1 <*> es2), + sumExprConst = sumExprConst e1 * sumExprConst e2 + } + +negateArithmeticalExpr :: ArithmeticalExpr -> ArithmeticalExpr +negateArithmeticalExpr (ArithmeticalExpr e) = ArithmeticalExpr $ negateSumExpr e + +plusArithmeticalExpr :: ArithmeticalExpr -> ArithmeticalExpr -> ArithmeticalExpr +plusArithmeticalExpr (ArithmeticalExpr e1) (ArithmeticalExpr e2) = ArithmeticalExpr $ plusSumExpr e1 e2 + +minusArithmeticalExpr :: ArithmeticalExpr -> ArithmeticalExpr -> ArithmeticalExpr +minusArithmeticalExpr (ArithmeticalExpr e1) (ArithmeticalExpr e2) = ArithmeticalExpr $ plusSumExpr e1 (negateSumExpr e2) + +multArithmeticalExpr :: ArithmeticalExpr -> ArithmeticalExpr -> ArithmeticalExpr +multArithmeticalExpr (ArithmeticalExpr e1) (ArithmeticalExpr e2) = ArithmeticalExpr $ multSumExpr e1 e2 + +parseSumExpr :: Expr -> SumExpr +parseSumExpr = \case + LitInt' n -> SumExpr {sumExprList = [], sumExprConst = n} + Negate' e -> negateSumExpr (parseSumExpr e) + Plus' e1 e2 -> plusSumExpr (parseSumExpr e1) (parseSumExpr e2) + Minus' e1 e2 -> plusSumExpr (parseSumExpr e1) (negateSumExpr (parseSumExpr e2)) + Mult' e1 e2 -> multSumExpr (parseSumExpr e1) (parseSumExpr e2) + e -> sumExprFromProductExpr (parseProductExpr e) + +-- | `parseArithmeticalExpr` converts a given expr to a normal form \(\sum_i \prod_j e _ {i,j})\). +-- This assumes given exprs have the type \(\mathbf{int}\). +parseArithmeticalExpr :: Expr -> ArithmeticalExpr +parseArithmeticalExpr = ArithmeticalExpr . parseSumExpr + +formatProductExpr :: ProductExpr -> Expr +formatProductExpr e = + let k = LitInt' (productExprConst e) + k' e' = case productExprConst e of + 0 -> LitInt' 0 + 1 -> e' + -1 -> Negate' e' + _ -> Mult' e' k + in case productExprList e of + [] -> k + eHead : esTail -> k' (foldl Mult' eHead esTail) + +formatSumExpr :: SumExpr -> Expr +formatSumExpr e = case sumExprList e of + [] -> LitInt' (sumExprConst e) + eHead : esTail -> + let op e' + | productExprConst e' > 0 = Plus' + | productExprConst e' < 0 = Minus' + | otherwise = const + go e1 e2 = op e2 e1 (formatProductExpr (e2 {productExprConst = abs (productExprConst e2)})) + k' e' + | sumExprConst e > 0 = Plus' e' (LitInt' (sumExprConst e)) + | sumExprConst e < 0 = Minus' e' (LitInt' (abs (sumExprConst e))) + | otherwise = e' + in k' (foldl go (formatProductExpr eHead) esTail) + +formatArithmeticalExpr :: ArithmeticalExpr -> Expr +formatArithmeticalExpr = formatSumExpr . unArithmeticalExpr . normalizeArithmeticalExpr + +normalizeProductExpr :: ProductExpr -> ProductExpr +normalizeProductExpr e = + let es = + if productExprConst e == 0 + then [] + else sort (productExprList e) + in e {productExprList = es} + +normalizeSumExpr :: SumExpr -> SumExpr +normalizeSumExpr e = + let cmp e1 e2 = productExprList e1 `compare` productExprList e2 + cmp' e1 e2 = cmp e1 e2 == EQ + es = sortBy cmp (map normalizeProductExpr (sumExprList e)) + es' = groupBy cmp' es + es'' = map (\group -> ProductExpr {productExprConst = sum (map productExprConst group), productExprList = productExprList (head group)}) es' + es''' = filter (\e -> productExprConst e /= 0 && not (null (productExprList e))) es'' + k = sum (map (\e -> if null (productExprList e) then productExprConst e else 0) es'') + in SumExpr + { sumExprList = es''', + sumExprConst = sumExprConst e + k + } + +normalizeArithmeticalExpr :: ArithmeticalExpr -> ArithmeticalExpr +normalizeArithmeticalExpr = ArithmeticalExpr . normalizeSumExpr . unArithmeticalExpr + +-- | `makeVectorFromArithmeticalExpr` makes a vector \(f\) and a expr \(c\) from a given vector of variables \(x_0, x_1, \dots, x _ {n - 1}\) and a given expr \(e\) s.t. \(f\) and \(c\) don't have \(x_0, x_1, \dots, x _ {n - 1}\) as free variables and \(e = c + f \cdot (x_0, x_1, \dots, x _ {n - 1})\) holds. +-- This assumes given variables and exprs have the type \(\mathbf{int}\). +-- +-- * The returned exprs are normalized with `normalizeArithmeticalExpr`. +makeVectorFromArithmeticalExpr :: V.Vector VarName -> ArithmeticalExpr -> Maybe (V.Vector ArithmeticalExpr, ArithmeticalExpr) +makeVectorFromArithmeticalExpr xs es = runST $ do + runMaybeT $ do + f <- lift $ MV.replicate (V.length xs) (integerArithmeticalExpr 0) + c <- lift $ newSTRef (integerArithmeticalExpr (sumExprConst (unArithmeticalExpr es))) + forM_ (sumExprList (unArithmeticalExpr es)) $ \e -> do + let indices = V.imap (\i x -> map (i,) (findIndices (x `isFreeVar`) (productExprList e))) xs + case concat (V.toList indices) of + [] -> lift $ modifySTRef c (plusArithmeticalExpr (arithmeticalExprFromProductExpr e)) + [(i, j)] -> do + let e' = e {productExprList = take j (productExprList e) ++ drop (j + 1) (productExprList e)} + lift $ MV.modify f (plusArithmeticalExpr (arithmeticalExprFromProductExpr e')) i + _ -> MaybeT $ return Nothing + f <- V.freeze f + c <- lift $ readSTRef c + return (V.map normalizeArithmeticalExpr f, normalizeArithmeticalExpr c) + +isZeroArithmeticalExpr :: ArithmeticalExpr -> Bool +isZeroArithmeticalExpr e = normalizeArithmeticalExpr e == integerArithmeticalExpr 0 + +isOneArithmeticalExpr :: ArithmeticalExpr -> Bool +isOneArithmeticalExpr e = normalizeArithmeticalExpr e == integerArithmeticalExpr 1 + +-- | `unNPlusKPattern` recognizes a pattern of \(x + k\) for a variable \(x\) and an integer constant \(k \in \mathbb{Z}\). +unNPlusKPattern :: ArithmeticalExpr -> Maybe (VarName, Integer) +unNPlusKPattern e = case normalizeArithmeticalExpr e of + ArithmeticalExpr + SumExpr + { sumExprList = + [ ProductExpr + { productExprConst = 1, + productExprList = [Var x] + } + ], + sumExprConst = k + } -> Just (x, k) + _ -> Nothing + +-- | `makeAffineFunctionFromArithmeticalExpr` is a specialized version of `makeVectorFromArithmeticalExpr`. +-- This function returns \(a, b\) for a given variable \(x\) and a given expr \(e = a x + b\) where \(a, b\) which doesn't use \(x\) free. +makeAffineFunctionFromArithmeticalExpr :: VarName -> ArithmeticalExpr -> Maybe (ArithmeticalExpr, ArithmeticalExpr) +makeAffineFunctionFromArithmeticalExpr x es = first V.head <$> makeVectorFromArithmeticalExpr (V.singleton x) es + +-- | `splitConstantFactorArithmeticalExpr` finds \(k\) and \(e'\) for given \(e\) s.t. \(e = k e'\). +splitConstantFactorArithmeticalExpr :: ArithmeticalExpr -> (Integer, ArithmeticalExpr) +splitConstantFactorArithmeticalExpr e = + let e' = unArithmeticalExpr $ normalizeArithmeticalExpr e + in case (sumExprConst e', sumExprList e') of + (0, []) -> (0, integerArithmeticalExpr 0) + (k, []) -> (k, integerArithmeticalExpr 1) + (0, [e]) -> second arithmeticalExprFromProductExpr $ splitConstantFactorProductExpr e + (k, es) -> + let kes = map splitConstantFactorProductExpr es + d = foldl gcd k (map fst kes) + in ( d, + ArithmeticalExpr + SumExpr + { sumExprConst = k `div` d, + sumExprList = map (\(k, e) -> e {productExprConst = (k * productExprConst e) `div` d}) kes + } + ) + +splitConstantFactorProductExpr :: ProductExpr -> (Integer, ProductExpr) +splitConstantFactorProductExpr e = (productExprConst e, e {productExprConst = 1}) diff --git a/src/Jikka/Core/Language/Beta.hs b/src/Jikka/Core/Language/Beta.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Language/Beta.hs @@ -0,0 +1,84 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Language.Beta +-- Description : does beta-reduction. / beta 簡約を行います。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Language.Beta + ( substitute, + substituteToplevelExpr, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Language.Expr +import Jikka.Core.Language.FreeVars +import Jikka.Core.Language.Util + +-- | `substitute` replaces the occrences of the given variable with the given expr. This considers contexts. +-- +-- >>> flip evalAlphaT 0 $ substitute (VarName "x") (Lit (LitInt 0)) (Lam (VarName "y") IntTy (Var (VarName "x"))) +-- Lam (VarName "y") IntTy (Lit (LitInt 0)) +-- +-- >>> flip evalAlphaT 0 $ substitute (VarName "x") (Lit (LitInt 0)) (Lam (VarName "x") IntTy (Var (VarName "x"))) +-- Lam (VarName "x") IntTy (Var (VarName "x")) +substitute :: MonadAlpha m => VarName -> Expr -> Expr -> m Expr +substitute x e = \case + Var y -> return $ if y == x then e else Var y + Lit lit -> return $ Lit lit + App e1 e2 -> App <$> substitute x e e1 <*> substitute x e e2 + Lam y t body -> + if x == y + then return $ Lam y t body + else do + (y, body) <- resolveConflict e (y, body) + Lam y t <$> substitute x e body + Let y t e1 e2 -> do + e1 <- substitute x e e1 + if y == x + then return $ Let y t e1 e2 + else do + (y, e2) <- resolveConflict e (y, e2) + Let y t e1 <$> substitute x e e2 + +substituteToplevelExpr :: (MonadAlpha m, MonadError Error m) => VarName -> Expr -> ToplevelExpr -> m ToplevelExpr +substituteToplevelExpr x e = \case + ResultExpr e' -> ResultExpr <$> substitute x e e' + ToplevelLet y t e' cont -> do + e' <- substitute x e e' + if y == x + then return $ ToplevelLet y t e' cont + else do + when (y `isFreeVar` e) $ do + throwInternalError $ "Jikka.Core.Language.Beta.substituteToplevelExpr: toplevel name conflicts: " ++ unVarName y + ToplevelLet y t e' <$> substituteToplevelExpr x e cont + ToplevelLetRec f args ret body cont -> do + if f == x + then return $ ToplevelLetRec f args ret body cont + else do + when (f `isFreeVar` e) $ do + throwInternalError $ "Jikka.Core.Language.Beta.substituteToplevelExpr: toplevel name conflicts: " ++ unVarName f + (args, body) <- + if x `elem` map fst args + then return (args, body) + else do + let go (args, body) (y, t) = do + (y, body) <- resolveConflict e (y, body) + return (args ++ [(y, t)], body) + foldM go ([], body) args + ToplevelLetRec f args ret body <$> substituteToplevelExpr x e cont + +resolveConflict :: MonadAlpha m => Expr -> (VarName, Expr) -> m (VarName, Expr) +resolveConflict e (x, e') = + if x `isFreeVar` e + then do + y <- genVarName x + e' <- substitute x (Var y) e' + return (y, e') + else return (x, e') diff --git a/src/Jikka/Core/Language/BuiltinPatterns.hs b/src/Jikka/Core/Language/BuiltinPatterns.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Language/BuiltinPatterns.hs @@ -0,0 +1,206 @@ +{-# LANGUAGE PatternSynonyms #-} + +-- | +-- Module : Jikka.Core.Language.BuiltinPatterns +-- Description : provides pattern synonyms for builtin functions. / 組み込み関数のための pattern synonyms を提供します。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- `Jikka.Core.Language.BuiltinPatterns` provides pattern synonyms for applications of `Builtin` functions. +-- For example, provide a pattern @Sum' e@ which is interpreted as @AppBuiltin Sum [e]@, or the same thing, @App (Lit (LitBuiltin Sum)) [e]@. +module Jikka.Core.Language.BuiltinPatterns where + +import Jikka.Core.Language.Expr + +-- arithmetical functions +pattern Negate' e = AppBuiltin Negate e + +pattern Plus' e1 e2 = AppBuiltin2 Plus e1 e2 + +pattern Minus' e1 e2 = AppBuiltin2 Minus e1 e2 + +pattern Mult' e1 e2 = AppBuiltin2 Mult e1 e2 + +pattern FloorDiv' e1 e2 = AppBuiltin2 FloorDiv e1 e2 + +pattern FloorMod' e1 e2 = AppBuiltin2 FloorMod e1 e2 + +pattern CeilDiv' e1 e2 = AppBuiltin2 CeilDiv e1 e2 + +pattern CeilMod' e1 e2 = AppBuiltin2 CeilMod e1 e2 + +pattern Pow' e1 e2 = AppBuiltin2 Pow e1 e2 + +-- advanced arithmetical functions +pattern Abs' e = AppBuiltin Abs e + +pattern Gcd' e1 e2 = AppBuiltin2 Gcd e1 e2 + +pattern Lcm' e1 e2 = AppBuiltin2 Lcm e1 e2 + +pattern Min2' t e1 e2 = AppBuiltin2 (Min2 t) e1 e2 + +pattern Max2' t e1 e2 = AppBuiltin2 (Max2 t) e1 e2 + +pattern Iterate' t n step base = AppBuiltin3 (Iterate t) n step base + +-- logical functions +pattern Not' e = AppBuiltin Not e + +pattern And' e1 e2 = AppBuiltin2 And e1 e2 + +pattern Or' e1 e2 = AppBuiltin2 Or e1 e2 + +pattern Implies' e1 e2 = AppBuiltin2 Implies e1 e2 + +pattern If' t e1 e2 e3 = AppBuiltin3 (If t) e1 e2 e3 + +-- bitwise functions +pattern BitNot' e = AppBuiltin BitNot e + +pattern BitAnd' e1 e2 = AppBuiltin2 BitAnd e1 e2 + +pattern BitOr' e1 e2 = AppBuiltin2 BitOr e1 e2 + +pattern BitXor' e1 e2 = AppBuiltin2 BitXor e1 e2 + +pattern BitLeftShift' e1 e2 = AppBuiltin2 BitLeftShift e1 e2 + +pattern BitRightShift' e1 e2 = AppBuiltin2 BitRightShift e1 e2 + +-- matrix functions + +pattern MatAp' h w e1 e2 = AppBuiltin2 (MatAp h w) e1 e2 + +pattern MatAdd' h w e1 e2 = AppBuiltin2 (MatAdd h w) e1 e2 + +pattern MatMul' h n w e1 e2 = AppBuiltin2 (MatMul h n w) e1 e2 + +pattern MatPow' n e1 e2 = AppBuiltin2 (MatPow n) e1 e2 + +pattern VecFloorMod' n e1 e2 = AppBuiltin2 (VecFloorMod n) e1 e2 + +pattern MatFloorMod' h w e1 e2 = AppBuiltin2 (MatFloorMod h w) e1 e2 + +-- modular functions +pattern ModNegate' e1 e2 = AppBuiltin2 ModNegate e1 e2 + +pattern ModPlus' e1 e2 e3 = AppBuiltin3 ModPlus e1 e2 e3 + +pattern ModMinus' e1 e2 e3 = AppBuiltin3 ModMinus e1 e2 e3 + +pattern ModMult' e1 e2 e3 = AppBuiltin3 ModMult e1 e2 e3 + +pattern ModInv' e1 e2 = AppBuiltin2 ModInv e1 e2 + +pattern ModPow' e1 e2 e3 = AppBuiltin3 ModPow e1 e2 e3 + +pattern ModMatAp' h w e1 e2 e3 = AppBuiltin3 (ModMatAp h w) e1 e2 e3 + +pattern ModMatAdd' h w e1 e2 e3 = AppBuiltin3 (ModMatAdd h w) e1 e2 e3 + +pattern ModMatMul' h n w e1 e2 e3 = AppBuiltin3 (ModMatMul h n w) e1 e2 e3 + +pattern ModMatPow' n e1 e2 e3 = AppBuiltin3 (ModMatPow n) e1 e2 e3 + +-- list functions +pattern Nil' t = Lit (LitNil t) + +pattern Cons' t e1 e2 = AppBuiltin2 (Cons t) e1 e2 + +pattern Snoc' t e1 e2 = AppBuiltin2 (Snoc t) e1 e2 + +pattern Foldl' t1 t2 e1 e2 e3 = AppBuiltin3 (Foldl t1 t2) e1 e2 e3 + +pattern Scanl' t1 t2 e1 e2 e3 = AppBuiltin3 (Scanl t1 t2) e1 e2 e3 + +pattern Build' t e1 e2 e3 = AppBuiltin3 (Build t) e1 e2 e3 + +pattern Len' t e = AppBuiltin (Len t) e + +pattern Map' t1 t2 f e = AppBuiltin2 (Map t1 t2) f e + +pattern Filter' t f e = AppBuiltin2 (Filter t) f e + +pattern At' t e1 e2 = AppBuiltin2 (At t) e1 e2 + +pattern SetAt' t e1 e2 e3 = AppBuiltin3 (SetAt t) e1 e2 e3 + +pattern Elem' t e1 e2 = AppBuiltin2 (Elem t) e1 e2 + +pattern Sum' e = AppBuiltin Sum e + +pattern Product' e = AppBuiltin Product e + +pattern ModSum' e1 e2 = AppBuiltin2 ModSum e1 e2 + +pattern ModProduct' e1 e2 = AppBuiltin2 ModProduct e1 e2 + +pattern Min1' t e = AppBuiltin (Min1 t) e + +pattern Max1' t e = AppBuiltin (Max1 t) e + +pattern ArgMin' t e = AppBuiltin (ArgMin t) e + +pattern ArgMax' t e = AppBuiltin (ArgMax t) e + +pattern All' e = AppBuiltin All e + +pattern Any' e = AppBuiltin Any e + +pattern Sorted' t e = AppBuiltin (Sorted t) e + +pattern Reversed' t e = AppBuiltin (Reversed t) e + +pattern Range1' e = AppBuiltin Range1 e + +pattern Range2' e1 e2 = AppBuiltin2 Range2 e1 e2 + +pattern Range3' e1 e2 e3 = AppBuiltin3 Range3 e1 e2 e3 + +-- tuple functions +pattern Tuple' ts = Lit (LitBuiltin (Tuple ts)) + +pattern Proj' ts n e = AppBuiltin (Proj ts n) e + +-- arithmetical relations +pattern LessThan' t e1 e2 = AppBuiltin2 (LessThan t) e1 e2 + +pattern LessEqual' t e1 e2 = AppBuiltin2 (LessEqual t) e1 e2 + +pattern GreaterThan' t e1 e2 = AppBuiltin2 (GreaterThan t) e1 e2 + +pattern GreaterEqual' t e1 e2 = AppBuiltin2 (GreaterEqual t) e1 e2 + +-- equality relations (polymorphic) +pattern Equal' t e1 e2 = AppBuiltin2 (Equal t) e1 e2 + +pattern NotEqual' t e1 e2 = AppBuiltin2 (NotEqual t) e1 e2 + +-- combinational functions +pattern Fact' e = AppBuiltin Fact e + +pattern Choose' e1 e2 = AppBuiltin2 Choose e1 e2 + +pattern Permute' e1 e2 = AppBuiltin2 Permute e1 e2 + +pattern MultiChoose' e1 e2 = AppBuiltin2 MultiChoose e1 e2 + +-- data structures +pattern ConvexHullTrickInit' = Lit (LitBuiltin ConvexHullTrickInit) + +pattern ConvexHullTrickGetMin' cht a = AppBuiltin2 ConvexHullTrickGetMin cht a + +pattern ConvexHullTrickInsert' cht a b = AppBuiltin3 ConvexHullTrickInsert cht a b + +pattern SegmentTreeInitList' semigrp a = AppBuiltin (SegmentTreeInitList semigrp) a + +pattern SegmentTreeGetRange' semigrp segtree e1 e2 = AppBuiltin3 (SegmentTreeGetRange semigrp) segtree e1 e2 + +pattern SegmentTreeSetPoint' semigrp segtree e1 e2 = AppBuiltin3 (SegmentTreeSetPoint semigrp) segtree e1 e2 + +-- errors +pattern Bottom' t err = Lit (LitBottom t err) diff --git a/src/Jikka/Core/Language/Expr.hs b/src/Jikka/Core/Language/Expr.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Language/Expr.hs @@ -0,0 +1,396 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE ViewPatterns #-} + +-- | +-- Module : Jikka.Core.Language.Expr +-- Description : has data types of our core language. / core 言語のためのデータ型を持ちます。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- `Jikka.Core.Language.Expr` module has the basic data types for our core language. +-- They are similar to the GHC Core language. +module Jikka.Core.Language.Expr where + +import Data.String (IsString) + +newtype VarName = VarName String deriving (Eq, Ord, Show, Read, IsString) + +unVarName :: VarName -> String +unVarName (VarName name) = name + +newtype TypeName = TypeName String deriving (Eq, Ord, Show, Read, IsString) + +unTypeName :: TypeName -> String +unTypeName (TypeName name) = name + +-- | `Type` represents the types of our core language. This is similar to the `Type` of GHC Core. +-- See also [commentary/compiler/type-type](https://gitlab.haskell.org/ghc/ghc/-/wikis/commentary/compiler/type-type). +-- +-- \[ +-- \newcommand\int{\mathbf{int}} +-- \newcommand\bool{\mathbf{bool}} +-- \newcommand\list{\mathbf{list}} +-- \begin{array}{rl} +-- \tau ::= & \alpha \\ +-- \vert & \int \\ +-- \vert & \bool \\ +-- \vert & \list(\tau) \\ +-- \vert & \tau \times \tau \times \dots \times \tau \\ +-- \vert & \tau \to \tau +-- \vert & \mathrm{data-structure} +-- \end{array} +-- \] +data Type + = VarTy TypeName + | IntTy + | BoolTy + | ListTy Type + | TupleTy [Type] + | FunTy Type Type + | DataStructureTy DataStructure + deriving (Eq, Ord, Show, Read) + +data DataStructure + = ConvexHullTrick + | SegmentTree Semigroup' + deriving (Eq, Ord, Show, Read) + +data Semigroup' + = SemigroupIntPlus + | SemigroupIntMin + | SemigroupIntMax + deriving (Eq, Ord, Show, Read) + +-- | TODO: What is the difference between `Literal` and `Builtin`? +data Builtin + = -- arithmetical functions + + -- | \(: \int \to \int\) + Negate + | -- | \(: \int \to \int \to \int\) + Plus + | -- | \(: \int \to \int \to \int\) + Minus + | -- | \(: \int \to \int \to \int\) + Mult + | -- | \(: \int \to \int \to \int\) + FloorDiv + | -- | \(: \int \to \int \to \int\) + FloorMod + | -- | \(: \int \to \int \to \int\) + CeilDiv + | -- | \(: \int \to \int \to \int\) + CeilMod + | -- | \(: \int \to \int \to \int\) + Pow + | -- advanced arithmetical functions + + -- | \(: \int \to \int\) + Abs + | -- | \(: \int \to \int \to \int\) + Gcd + | -- | \(: \int \to \int \to \int\) + Lcm + | -- | \(: \forall \alpha. \alpha \to \alpha \to \alpha\) + Min2 Type + | -- | \(: \forall \alpha. \alpha \to \alpha \to \alpha\) + Max2 Type + | -- | iterated application \((\lambda k f x. f^k(x)): \forall \alpha. \int \to (\alpha \to \alpha) \to \alpha \to \alpha\) + Iterate Type + | -- logical functions + + -- | \(: \bool \to \bool\) + Not + | -- | \(: \bool \to \bool \to \bool\) + And + | -- | \(: \bool \to \bool \to \bool\) + Or + | -- | \(: \bool \to \bool \to \bool\) + Implies + | -- | \(: \forall \alpha. \bool \to \alpha \to \alpha \to \alpha\) + If Type + | -- bitwise functions + + -- | \(: \int \to \int\) + BitNot + | -- | \(: \int \to \int \to \int\) + BitAnd + | -- | \(: \int \to \int \to \int\) + BitOr + | -- | \(: \int \to \int \to \int\) + BitXor + | -- | \(: \int \to \int \to \int\) + BitLeftShift + | -- | \(: \int \to \int \to \int\) + BitRightShift + | -- matrix functions + + -- | matrix application \(: \int^{H \times W} \to \int^W \to \int^H\) + MatAp Int Int + | -- | zero matrix \(: \to \int^{n \times n}\) + MatZero Int + | -- | unit matrix \(: \to \int^{n \times n}\) + MatOne Int + | -- | matrix addition \(: \int^{H \times W} \to \int^{H \times W} \to \int^{H \times W}\) + MatAdd Int Int + | -- | matrix multiplication \(: \int^{H \times n} \to \int^{n \times W} \to \int^{H \times W}\) + MatMul Int Int Int + | -- | matrix power \(: \int^{n \times n} \to \int \to \int^{n \times n}\) + MatPow Int + | -- | vector point-wise floor-mod \(: \int^{n} \to \int \to \int^{n}\) + VecFloorMod Int + | -- | matrix point-wise floor-mod \(: \int^{H \times W} \to \int \to \int^{H \times W}\) + MatFloorMod Int Int + | -- modular functions + + -- | \(: \int \to \int \to \int\) + ModNegate + | -- | \(: \int \to \int \to \int \to \int\) + ModPlus + | -- | \(: \int \to \int \to \int \to \int\) + ModMinus + | -- | \(: \int \to \int \to \int \to \int\) + ModMult + | -- | \(: \int \to \int \to \int\) + ModInv + | -- | \(: \int \to \int \to \int \to \int\) + ModPow + | -- | matrix application \(: \int^{H \times W} \to \int^W \to \int \to \int^H\) + ModMatAp Int Int + | -- | matrix addition \(: \int^{H \times W} \to \int^{H \times W} \to \int \to \int^{H \times W}\) + ModMatAdd Int Int + | -- | matrix multiplication \(: \int^{H \times n} \to \int^{n \times W} \to \int \to \int^{H \times W}\) + ModMatMul Int Int Int + | -- | matrix power \(: \int^{n \times n} \to \int \to \int^{n \times n}\) + ModMatPow Int + | -- list functions + + -- | \(: \forall \alpha. \alpha \to \list(\alpha) \to \list(\alpha)\) + Cons Type + | -- | \(: \forall \alpha. \list(alpha) \to \alpha \to \list(\alpha)\) + Snoc Type + | -- | \(: \forall \alpha \beta. (\beta \to \alpha \to \beta) \to \beta \to \list(\alpha) \to \beta\) + Foldl Type Type + | -- | \(: \forall \alpha \beta. (\beta \to \alpha \to \beta) \to \beta \to \list(\alpha) \to \list(\beta)\) + Scanl Type Type + | -- | \(\lambda f a n.\) repeat @a <- snoc a (f a)@ @n@ times \(: \forall \alpha. (\list(\alpha) \to \alpha) \to \list(\alpha) \to \int \to \list(\alpha)\) + Build Type + | -- | \(: \forall \alpha. \list(\alpha) \to \int\) + Len Type + | -- | \(: \forall \alpha \beta. (\alpha \to \beta) \to \list(\alpha) \to \list(\beta)\) + Map Type Type + | -- | \(: \forall \alpha \beta. (\alpha \to \bool) \to \list(\alpha) \to \list(\beta)\) + Filter Type + | -- | \(: \forall \alpha. \list(\alpha) \to \int \to \alpha\) + At Type + | -- | \(: \forall \alpha. \list(\alpha) \to \int \to \alpha \to \list(\alpha)\) + SetAt Type + | -- | \(: \forall \alpha. \alpha \to \list(\alpha) \to \bool\) + Elem Type + | -- | \(: \list(\int) \to \int\) + Sum + | -- | \(: \list(\int) \to \int\) + Product + | -- | \(: \list(\int) \to \int \to \int\) + ModSum + | -- | \(: \list(\int) \to \int \to \int\) + ModProduct + | -- | \(: \forall \alpha. \list(\alpha) \to \alpha\) + Min1 Type + | -- | \(: \forall \alpha. \list(\alpha) \to \alpha\) + Max1 Type + | -- | \(: \forall \alpha. \list(\alpha) \to \int\) + ArgMin Type + | -- | \(: \forall \alpha. \list(\alpha) \to \int\) + ArgMax Type + | -- | \(: \list(\bool) \to \bool\) + All + | -- | \(: \list(\bool) \to \bool\) + Any + | -- | \(: \forall \alpha. \list(\alpha) \to \list(\alpha)\) + Sorted Type + | -- | \(: \forall \alpha. \list(\alpha) \to \list(\alpha)\) + Reversed Type + | -- | \(: \int \to \list(\int)\) + Range1 + | -- | \(: \int \to \int \to \list(\int)\) + Range2 + | -- | \(: \int \to \int \to \int \to \list(\int)\) + Range3 + | -- tuple functions + + -- | \(: \forall \alpha_0 \alpha_1 \dots \alpha _ {n - 1}. \alpha_0 \to \dots \to \alpha _ {n - 1} \to \alpha_0 \times \dots \times \alpha _ {n - 1}\) + Tuple [Type] + | -- | \(: \forall \alpha_0 \alpha_1 \dots \alpha _ {n - 1}. \alpha_0 \times \dots \times \alpha _ {n - 1} \to \alpha_i\) + Proj [Type] Int + | -- comparison + + -- | \(: \forall \alpha. \alpha \to \alpha \to \bool\) + LessThan Type + | -- | \(: \forall \alpha. \alpha \to \alpha \to \bool\) + LessEqual Type + | -- | \(: \forall \alpha. \alpha \to \alpha \to \bool\) + GreaterThan Type + | -- | \(: \forall \alpha. \alpha \to \alpha \to \bool\) + GreaterEqual Type + | -- | \(: \forall \alpha. \alpha \to \alpha \to \bool\) + Equal Type + | -- | \(: \forall \alpha. \alpha \to \alpha \to \bool\) + NotEqual Type + | -- combinational functions + + -- | \(: \int \to \int\) + Fact + | -- | \(: \int \to \int \to \int\) + Choose + | -- | \(: \int \to \int \to \int\) + Permute + | -- | \(: \int \to \int \to \int\) + MultiChoose + | -- data structures + + -- | \(: \mathrm{convex-hull-trick}\) + ConvexHullTrickInit + | -- | \(: \mathrm{convex-hull-trick} \to \int \to \int\) + ConvexHullTrickGetMin + | -- | \(: \mathrm{convex-hull-trick} \to \int \to \int \to \mathrm{convex-hull-trick}\) + ConvexHullTrickInsert + | -- | \(: \forall S. \list(S) \to \mathrm{segment-tree}(S)\) + SegmentTreeInitList Semigroup' + | -- | \(: \forall S. \mathrm{segment-tree}(S) \to \int \to \int \to S\) + SegmentTreeGetRange Semigroup' + | -- | \(: \forall S. \mathrm{segment-tree}(S) \to \int \to S \to \mathrm{segment-tree}(S)\) + SegmentTreeSetPoint Semigroup' + deriving (Eq, Ord, Show, Read) + +data Literal + = LitBuiltin Builtin + | -- | \(: \forall \alpha. \int\) + LitInt Integer + | -- | \(: \forall \alpha. \bool\) + LitBool Bool + | -- | \(: \forall \alpha. \list(\alpha)\) + LitNil Type + | -- | \(: \bot : \forall \alpha. \alpha\). The second argument is its error message. + LitBottom Type String + deriving (Eq, Ord, Show, Read) + +-- | `Expr` represents the exprs of our core language. This is similar to the `Expr` of GHC Core. +-- See also [commentary/compiler/core-syn-type](https://gitlab.haskell.org/ghc/ghc/-/wikis/commentary/compiler/core-syn-type). +-- +-- \[ +-- \begin{array}{rl} +-- e ::= & x \\ +-- \vert & \mathrm{literal}\ldots \\ +-- \vert & e_0(e_1, e_2, \dots, e_n) \\ +-- \vert & \lambda ~ x_0\colon \tau_0, x_1\colon \tau_1, \dots, x_{n-1}\colon \tau_{n-1}. ~ e \\ +-- \vert & \mathbf{let} ~ x\colon \tau = e_1 ~ \mathbf{in} ~ e_2 +-- \end{array} +-- \] +data Expr + = Var VarName + | Lit Literal + | -- | The functions are not curried. + App Expr Expr + | -- | The lambdas are also not curried. + Lam VarName Type Expr + | -- | This "let" is not recursive. + Let VarName Type Expr Expr + deriving (Eq, Ord, Show, Read) + +pattern Fun2Ty t1 t2 ret = FunTy t1 (FunTy t2 ret) + +pattern Fun3Ty t1 t2 t3 ret = FunTy t1 (FunTy t2 (FunTy t3 ret)) + +pattern Fun1STy t <- + (\case FunTy t1 ret | t1 == ret -> Just ret; _ -> Nothing -> Just t) + where + Fun1STy t = FunTy t t + +pattern Fun2STy t <- + (\case Fun2Ty t1 t2 ret | t1 == ret && t2 == ret -> Just ret; _ -> Nothing -> Just t) + where + Fun2STy t = Fun2Ty t t t + +pattern Fun3STy t <- + (\case Fun3Ty t1 t2 t3 ret | t1 == ret && t2 == ret && t3 == ret -> Just ret; _ -> Nothing -> Just t) + where + Fun3STy t = Fun3Ty t t t t + +pattern FunLTy t <- + (\case FunTy (ListTy t1) ret | t1 == ret -> Just ret; _ -> Nothing -> Just t) + where + FunLTy t = FunTy (ListTy t) t + +vectorTy :: Int -> Type +vectorTy n = TupleTy (replicate n IntTy) + +matrixTy :: Int -> Int -> Type +matrixTy h w = TupleTy (replicate h (TupleTy (replicate w IntTy))) + +pattern UnitTy = TupleTy [] + +pattern ConvexHullTrickTy = DataStructureTy ConvexHullTrick + +pattern SegmentTreeTy semigrp = DataStructureTy (SegmentTree semigrp) + +pattern LitInt' n = Lit (LitInt n) + +pattern Lit0 = Lit (LitInt 0) + +pattern Lit1 = Lit (LitInt 1) + +pattern Lit2 = Lit (LitInt 2) + +pattern LitMinus1 = Lit (LitInt (-1)) + +pattern LitBool' p = Lit (LitBool p) + +pattern LitTrue = Lit (LitBool True) + +pattern LitFalse = Lit (LitBool False) + +pattern Builtin builtin = Lit (LitBuiltin builtin) + +pattern App2 f e1 e2 = App (App f e1) e2 + +pattern App3 f e1 e2 e3 = App (App (App f e1) e2) e3 + +pattern App4 f e1 e2 e3 e4 = App (App (App (App f e1) e2) e3) e4 + +pattern AppBuiltin builtin e1 = App (Lit (LitBuiltin builtin)) e1 + +pattern AppBuiltin2 builtin e1 e2 = App2 (Lit (LitBuiltin builtin)) e1 e2 + +pattern AppBuiltin3 builtin e1 e2 e3 = App3 (Lit (LitBuiltin builtin)) e1 e2 e3 + +pattern Lam2 x1 t1 x2 t2 e = Lam x1 t1 (Lam x2 t2 e) + +pattern Lam3 x1 t1 x2 t2 x3 t3 e = Lam x1 t1 (Lam x2 t2 (Lam x3 t3 e)) + +pattern LamId x t <- + (\case Lam x t (Var y) | x == y -> Just (x, t); _ -> Nothing -> Just (x, t)) + where + LamId x t = Lam x t (Var x) + +-- | `ToplevelExpr` is the toplevel exprs. In our core, "let rec" is allowed only on the toplevel. +-- +-- \[ +-- \begin{array}{rl} +-- \mathrm{tle} ::= & e \\ +-- \vert & \mathbf{let}~ x: \tau = e ~\mathbf{in}~ \mathrm{tle} \\ +-- \vert & \mathbf{let~rec}~ x(x: \tau, x: \tau, \dots, x: \tau): \tau = e ~\mathbf{in}~ \mathrm{tle} +-- \end{array} +-- \] +data ToplevelExpr + = ResultExpr Expr + | ToplevelLet VarName Type Expr ToplevelExpr + | ToplevelLetRec VarName [(VarName, Type)] Type Expr ToplevelExpr + deriving (Eq, Ord, Show, Read) + +type Program = ToplevelExpr diff --git a/src/Jikka/Core/Language/FreeVars.hs b/src/Jikka/Core/Language/FreeVars.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Language/FreeVars.hs @@ -0,0 +1,62 @@ +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Converter.Core.FreeVars +-- Description : provides utilities aboud free variables. / 自由変数についてのユーティリティを提供します。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Language.FreeVars where + +import Jikka.Core.Language.Expr + +-- | `isFreeVar` checks if the given variable occurs in the tiven expr. This considers contexts. +-- +-- >>> VarName "x" `isFreeVar` Lam (VarName "y") IntTy (Var (VarName "x")) +-- True +-- +-- >>> VarName "x" `isFreeVar` Lam (VarName "x") IntTy (Var (VarName "x")) +-- False +isFreeVar :: VarName -> Expr -> Bool +isFreeVar x = \case + Var y -> y == x + Lit _ -> False + App f e -> isFreeVar x f || isFreeVar x e + Lam y _ e -> x /= y && isFreeVar x e + Let y _ e1 e2 -> (y /= x && isFreeVar x e1) || isFreeVar x e2 + +-- | `isUnusedVar` is the negation of `isFreeVar`. +-- +-- TODO: rename to `isNonFreeVar`? +isUnusedVar :: VarName -> Expr -> Bool +isUnusedVar x e = not (isFreeVar x e) + +-- | `isFreeVarOrScopedVar` checks if the given variable occurs in the tiven expr. This ignores contexts. +-- +-- >>> VarName "x" `isFreeVarOrScopedVar` Lam (VarName "x") IntTy (Var (VarName "y")) +-- True +isFreeVarOrScopedVar :: VarName -> Expr -> Bool +isFreeVarOrScopedVar x = \case + Var y -> y == x + Lit _ -> False + App f e -> isFreeVarOrScopedVar x f || isFreeVarOrScopedVar x e + Lam y _ e -> x == y || isFreeVarOrScopedVar x e + Let y _ e1 e2 -> y == x || isFreeVarOrScopedVar x e1 || isFreeVarOrScopedVar x e2 + +freeTyVars :: Type -> [TypeName] +freeTyVars = \case + VarTy x -> [x] + IntTy -> [] + BoolTy -> [] + ListTy t -> freeTyVars t + TupleTy ts -> concatMap freeTyVars ts + FunTy t1 t2 -> freeTyVars t1 ++ freeTyVars t2 + DataStructureTy _ -> [] + +findUnusedVarName :: VarName -> Expr -> VarName +findUnusedVarName (VarName x) e = head . filter (`isUnusedVar` e) $ map (\i -> VarName (x ++ show i)) [0 ..] + +findUnusedVarName' :: Expr -> VarName +findUnusedVarName' = findUnusedVarName (VarName "x") diff --git a/src/Jikka/Core/Language/Lint.hs b/src/Jikka/Core/Language/Lint.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Language/Lint.hs @@ -0,0 +1,33 @@ +{-# LANGUAGE FlexibleContexts #-} + +-- | +-- Module : Jikka.Core.Language.Lint +-- Description : verifies various conditions (e.g. well-typed) of exprs of our core language. / core 言語の式の種々の条件 (例: 型付け可能性) を検査します。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- `Jikka.Core.Language.Lint` module checks the invariants of data types. Mainly, this checks types of `Expr`. +module Jikka.Core.Language.Lint where + +import Jikka.Common.Error +import Jikka.Core.Language.Expr +import Jikka.Core.Language.TypeCheck + +precondition :: MonadError Error m => m a -> m a +precondition = wrapError' "precondition" + +postcondition :: MonadError Error m => m a -> m a +postcondition = wrapError' "postcondition" + +-- | TODO: implement this +ensureEagerlyEvaluatable :: MonadError Error m => Program -> m () +ensureEagerlyEvaluatable _ = wrapError' "Jikka.Core.Language.Lint.ensureEagerlyEvaluatable" $ do + return () + +ensureWellTyped :: MonadError Error m => Program -> m () +ensureWellTyped prog = wrapError' "Jikka.Core.Language.Lint.ensureWellTyped" $ do + _ <- typecheckProgram prog + return () diff --git a/src/Jikka/Core/Language/RewriteRules.hs b/src/Jikka/Core/Language/RewriteRules.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Language/RewriteRules.hs @@ -0,0 +1,141 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Language.RewriteRules +-- Description : provides functions for rewrite rules. / 書き換え規則のための関数を提供します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Language.RewriteRules + ( RewriteRule (..), + pureRewriteRule, + simpleRewriteRule, + applyRewriteRule, + applyRewriteRuleToplevelExpr, + applyRewriteRuleProgram, + applyRewriteRuleProgram', + traceRewriteRule, + ) +where + +import Control.Monad.State.Strict +import Data.Maybe (fromMaybe) +import Debug.Trace +import Jikka.Common.Error +import Jikka.Core.Format (formatExpr) +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Util (curryFunTy) + +newtype RewriteRule m = RewriteRule ([(VarName, Type)] -> Expr -> m (Maybe Expr)) + +unRewriteRule :: RewriteRule m -> [(VarName, Type)] -> Expr -> m (Maybe Expr) +unRewriteRule (RewriteRule f) = f + +instance Monad m => Semigroup (RewriteRule m) where + f <> g = RewriteRule $ \env e -> do + e' <- unRewriteRule f env e + case e' of + Nothing -> unRewriteRule g env e + Just e' -> do + e'' <- unRewriteRule g env e' + case e'' of + Nothing -> return (Just e') + Just e'' -> return (Just e'') + +instance Monad m => Monoid (RewriteRule m) where + mempty = RewriteRule (\_ _ -> return Nothing) + +pureRewriteRule :: Monad m => ([(VarName, Type)] -> Expr -> Maybe Expr) -> RewriteRule m +pureRewriteRule f = RewriteRule (\env e -> return (f env e)) + +simpleRewriteRule :: Monad m => (Expr -> Maybe Expr) -> RewriteRule m +simpleRewriteRule f = RewriteRule (\_ e -> return (f e)) + +-- | `applyRewriteRule` applies a given rule to a given expr. +-- This rewrites on all sub-exprs of the given expr, and repeats to rewrite while it is possible. +-- +-- * This function is idempotent. +-- * This function doesn't terminate when a given rewrite rule doesn't terminate. +applyRewriteRule :: MonadError Error m => RewriteRule m -> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr) +applyRewriteRule = applyRewriteRulePreOrder + +coalesceMaybes :: a -> Maybe a -> b -> Maybe b -> Maybe (a, b) +coalesceMaybes _ Nothing _ Nothing = Nothing +coalesceMaybes a Nothing _ (Just b) = Just (a, b) +coalesceMaybes _ (Just a) b Nothing = Just (a, b) +coalesceMaybes _ (Just a) _ (Just b) = Just (a, b) + +applyRewriteRuleToImmediateSubExprs :: MonadError Error m => RewriteRule m -> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr) +applyRewriteRuleToImmediateSubExprs f env = \case + Var _ -> return Nothing + Lit _ -> return Nothing + App e1 e2 -> do + e1' <- lift $ unRewriteRule f env e1 + e2' <- lift $ unRewriteRule f env e2 + return $ fmap (uncurry App) (coalesceMaybes e1 e1' e2 e2') + Lam x t body -> lift $ (Lam x t <$>) <$> unRewriteRule f ((x, t) : env) body + Let x t e1 e2 -> do + e1' <- lift $ unRewriteRule f env e1 + e2' <- lift $ unRewriteRule f ((x, t) : env) e2 + return $ fmap (uncurry (Let x t)) (coalesceMaybes e1 e1' e2 e2') + +joinStateT :: Monad m => StateT s (StateT s m) a -> StateT s m a +joinStateT f = do + s <- get + (a, s) <- runStateT f s + put s + return a + +applyRewriteRulePreOrder :: MonadError Error m => RewriteRule m -> [(VarName, Type)] -> Expr -> StateT Integer m (Maybe Expr) +applyRewriteRulePreOrder f env e = do + cnt <- get + when (cnt >= 100) $ do + throwInternalError "rewrite rule doesn't terminate" + e' <- lift $ unRewriteRule f env e + case e' of + Nothing -> do + e' <- joinStateT (applyRewriteRuleToImmediateSubExprs (RewriteRule (applyRewriteRulePreOrder f)) env e) + case e' of + Nothing -> return Nothing + Just e' -> do + modify' succ + e'' <- lift $ unRewriteRule f env e' + case e'' of + Nothing -> return $ Just e' + Just e'' -> do + modify' succ + e''' <- applyRewriteRulePreOrder f env e'' + return . Just $ fromMaybe e'' e''' + Just e' -> do + modify' succ + e'' <- applyRewriteRulePreOrder f env e' + return . Just $ fromMaybe e' e'' + +applyRewriteRuleToplevelExpr :: MonadError Error m => RewriteRule m -> [(VarName, Type)] -> ToplevelExpr -> StateT Integer m (Maybe ToplevelExpr) +applyRewriteRuleToplevelExpr f env = \case + ResultExpr e -> (ResultExpr <$>) <$> applyRewriteRule f env e + ToplevelLet y t e cont -> do + e' <- applyRewriteRule f env e + cont' <- applyRewriteRuleToplevelExpr f ((y, t) : env) cont + return $ fmap (uncurry (ToplevelLet y t)) (coalesceMaybes e e' cont cont') + ToplevelLetRec g args ret body cont -> do + let env' = (g, curryFunTy (map snd args) ret) : env + body' <- applyRewriteRule f (reverse args ++ env') body + cont' <- applyRewriteRuleToplevelExpr f env' cont + return $ fmap (uncurry (ToplevelLetRec g args ret)) (coalesceMaybes body body' cont cont') + +applyRewriteRuleProgram :: MonadError Error m => RewriteRule m -> Program -> m (Maybe Program) +applyRewriteRuleProgram f prog = evalStateT (applyRewriteRuleToplevelExpr f [] prog) 0 + +applyRewriteRuleProgram' :: MonadError Error m => RewriteRule m -> Program -> m Program +applyRewriteRuleProgram' f prog = fromMaybe prog <$> applyRewriteRuleProgram f prog + +traceRewriteRule :: Monad m => RewriteRule m -> RewriteRule m +traceRewriteRule f = RewriteRule $ \env e -> do + e' <- unRewriteRule f env e + case e' of + Nothing -> return Nothing + Just e' -> trace ("before:\n" ++ formatExpr e ++ "\nafter:\n" ++ formatExpr e') (return (Just e')) diff --git a/src/Jikka/Core/Language/Runtime.hs b/src/Jikka/Core/Language/Runtime.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Language/Runtime.hs @@ -0,0 +1,52 @@ +{-# LANGUAGE FlexibleContexts #-} + +module Jikka.Core.Language.Runtime where + +import Jikka.Common.Error + +floorDiv :: MonadError Error m => Integer -> Integer -> m Integer +floorDiv _ 0 = throwRuntimeError "zero div" +floorDiv a b = return (a `div` b) + +floorMod :: MonadError Error m => Integer -> Integer -> m Integer +floorMod _ 0 = throwRuntimeError "zero div" +floorMod a b = return (a `mod` b) + +ceilDiv :: MonadError Error m => Integer -> Integer -> m Integer +ceilDiv _ 0 = throwRuntimeError "zero div" +ceilDiv a b = return ((a + b - 1) `div` b) + +ceilMod :: MonadError Error m => Integer -> Integer -> m Integer +ceilMod _ 0 = throwRuntimeError "zero div" +ceilMod a b = return (a - ((a + b - 1) `div` b) * b) + +modinv :: MonadError Error m => Integer -> Integer -> m Integer +modinv a m | m <= 0 || a `mod` m == 0 = throwRuntimeError $ "invalid argument for inv: " ++ show (a, m) +modinv a m = go a m 0 1 1 0 + where + go 0 b x y _ _ = if a * x + m * y == b && b == 1 then return x else throwRuntimeError "Jikka.Core.Language.Runtime.modinv: something wrong" + go a b x y u v = let q = b `div` a in go (b - q * a) a u v (x - q * u) (y - q * v) + +modpow :: MonadError Error m => Integer -> Integer -> Integer -> m Integer +modpow _ _ m | m <= 0 = throwRuntimeError $ "invalid argument for modpow: MOD = " ++ show m +modpow a b m = return $ go (a `mod` m) b + where + go a 0 = a + go a b = go (if b `mod` 2 == 1 then a * b `mod` m else a) (b `div` 2) + +fact :: MonadError Error m => Integer -> m Integer +fact n | n < 0 = throwRuntimeError $ "invalid argument for fact: " ++ show n +fact n = return $ product [1 .. n] + +choose :: MonadError Error m => Integer -> Integer -> m Integer +choose n r | not (0 <= r && r <= n) = throwRuntimeError $ "invalid argument for choose: " ++ show (n, r) +choose n r = return $ product [n - r + 1 .. n] `div` product [1 .. r] + +permute :: MonadError Error m => Integer -> Integer -> m Integer +permute n r | not (0 <= r && r <= n) = throwRuntimeError $ "invalid argument for choose: " ++ show (n, r) +permute n r = return $ product [n - r + 1 .. n] + +multichoose :: MonadError Error m => Integer -> Integer -> m Integer +multichoose n r | not (0 <= r && r <= n) = throwRuntimeError $ "invalid argument for multichoose: " ++ show (n, r) +multichoose 0 0 = return 1 +multichoose n r = choose (n + r - 1) r diff --git a/src/Jikka/Core/Language/TypeCheck.hs b/src/Jikka/Core/Language/TypeCheck.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Language/TypeCheck.hs @@ -0,0 +1,188 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Core.Language.RewriteRules +-- Description : checks and obtains types of exprs. / 式の型を検査し取得します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Core.Language.TypeCheck where + +import Jikka.Common.Error +import Jikka.Core.Format (formatExpr, formatType) +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Util + +builtinToType :: Builtin -> Type +builtinToType = \case + -- arithmetical functions + Negate -> Fun1STy IntTy + Plus -> Fun2STy IntTy + Minus -> Fun2STy IntTy + Mult -> Fun2STy IntTy + FloorDiv -> Fun2STy IntTy + FloorMod -> Fun2STy IntTy + CeilDiv -> Fun2STy IntTy + CeilMod -> Fun2STy IntTy + Pow -> Fun2STy IntTy + -- advanced arithmetical functions + Abs -> Fun1STy IntTy + Gcd -> Fun2STy IntTy + Lcm -> Fun2STy IntTy + Min2 t -> Fun2STy t + Max2 t -> Fun2STy t + Iterate t -> Fun3Ty IntTy (FunTy t t) t t + -- logical functions + Not -> Fun1STy BoolTy + And -> Fun2STy BoolTy + Or -> Fun2STy BoolTy + Implies -> Fun2STy BoolTy + If t -> Fun3Ty BoolTy t t t + -- bitwise functions + BitNot -> Fun1STy IntTy + BitAnd -> Fun2STy IntTy + BitOr -> Fun2STy IntTy + BitXor -> Fun2STy IntTy + BitLeftShift -> Fun2STy IntTy + BitRightShift -> Fun2STy IntTy + -- matrix functions + MatAp h w -> Fun2Ty (matrixTy h w) (vectorTy w) (vectorTy h) + MatZero n -> matrixTy n n + MatOne n -> matrixTy n n + MatAdd h w -> Fun2Ty (matrixTy h w) (matrixTy h w) (matrixTy h w) + MatMul h n w -> Fun2Ty (matrixTy h n) (matrixTy n w) (matrixTy h w) + MatPow n -> Fun2Ty (matrixTy n n) IntTy (matrixTy n n) + VecFloorMod n -> Fun2Ty (vectorTy n) IntTy (vectorTy n) + MatFloorMod h w -> Fun2Ty (matrixTy h w) IntTy (matrixTy h w) + -- modular functions + ModNegate -> Fun2STy IntTy + ModPlus -> Fun3STy IntTy + ModMinus -> Fun3STy IntTy + ModMult -> Fun3STy IntTy + ModInv -> Fun2STy IntTy + ModPow -> Fun3STy IntTy + ModMatAp h w -> Fun3Ty (matrixTy h w) (vectorTy w) IntTy (vectorTy h) + ModMatAdd h w -> Fun3Ty (matrixTy h w) (matrixTy h w) IntTy (matrixTy h w) + ModMatMul h n w -> Fun3Ty (matrixTy h n) (matrixTy n w) IntTy (matrixTy h w) + ModMatPow n -> Fun3Ty (matrixTy n n) IntTy IntTy (matrixTy n n) + -- list functions + Cons t -> Fun2Ty t (ListTy t) (ListTy t) + Snoc t -> Fun2Ty (ListTy t) t (ListTy t) + Foldl t1 t2 -> Fun3Ty (Fun2Ty t2 t1 t2) t2 (ListTy t1) t2 + Scanl t1 t2 -> Fun3Ty (Fun2Ty t2 t1 t2) t2 (ListTy t1) (ListTy t2) + Build t -> Fun3Ty (FunTy (ListTy t) t) (ListTy t) IntTy (ListTy t) + Len t -> FunTy (ListTy t) IntTy + Map t1 t2 -> Fun2Ty (FunTy t1 t2) (ListTy t1) (ListTy t2) + Filter t -> Fun2Ty (FunTy t BoolTy) (ListTy t) (ListTy t) + At t -> Fun2Ty (ListTy t) IntTy t + SetAt t -> Fun3Ty (ListTy t) IntTy t (ListTy t) + Elem t -> Fun2Ty t (ListTy t) BoolTy + Sum -> FunLTy IntTy + Product -> FunLTy IntTy + ModSum -> Fun2Ty (ListTy IntTy) IntTy IntTy + ModProduct -> Fun2Ty (ListTy IntTy) IntTy IntTy + Min1 t -> FunLTy t + Max1 t -> FunLTy t + ArgMin t -> FunTy (ListTy t) IntTy + ArgMax t -> FunTy (ListTy t) IntTy + All -> FunLTy BoolTy + Any -> FunLTy BoolTy + Sorted t -> Fun1STy (ListTy t) + Reversed t -> Fun1STy (ListTy t) + Range1 -> FunTy IntTy (ListTy IntTy) + Range2 -> Fun2Ty IntTy IntTy (ListTy IntTy) + Range3 -> Fun3Ty IntTy IntTy IntTy (ListTy IntTy) + -- tuple functions + Tuple ts -> curryFunTy ts (TupleTy ts) + Proj ts n -> FunTy (TupleTy ts) (ts !! n) + -- comparison + LessThan t -> Fun2Ty t t BoolTy + LessEqual t -> Fun2Ty t t BoolTy + GreaterThan t -> Fun2Ty t t BoolTy + GreaterEqual t -> Fun2Ty t t BoolTy + Equal t -> Fun2Ty t t BoolTy + NotEqual t -> Fun2Ty t t BoolTy + -- combinational functions + Fact -> Fun1STy IntTy + Choose -> Fun2STy IntTy + Permute -> Fun2STy IntTy + MultiChoose -> Fun2STy IntTy + -- data structure + ConvexHullTrickInit -> ConvexHullTrickTy + ConvexHullTrickGetMin -> Fun2Ty ConvexHullTrickTy IntTy IntTy + ConvexHullTrickInsert -> Fun3Ty ConvexHullTrickTy IntTy IntTy ConvexHullTrickTy + SegmentTreeInitList semigrp -> FunTy (ListTy (semigroupToType semigrp)) (SegmentTreeTy semigrp) + SegmentTreeGetRange semigrp -> Fun3Ty (SegmentTreeTy semigrp) IntTy IntTy (semigroupToType semigrp) + SegmentTreeSetPoint semigrp -> Fun3Ty (SegmentTreeTy semigrp) IntTy (semigroupToType semigrp) (SegmentTreeTy semigrp) + +semigroupToType :: Semigroup' -> Type +semigroupToType = \case + SemigroupIntPlus -> IntTy + SemigroupIntMin -> IntTy + SemigroupIntMax -> IntTy + +literalToType :: Literal -> Type +literalToType = \case + LitBuiltin builtin -> builtinToType builtin + LitInt _ -> IntTy + LitBool _ -> BoolTy + LitNil t -> ListTy t + LitBottom t _ -> t + +arityOfBuiltin :: Builtin -> Int +arityOfBuiltin = \case + Min2 _ -> 2 + Max2 _ -> 2 + Foldl _ _ -> 3 + Iterate _ -> 3 + At _ -> 2 + Min1 _ -> 1 + Max1 _ -> 1 + Proj _ _ -> 1 + builtin -> length (fst (uncurryFunTy (builtinToType builtin))) + +type TypeEnv = [(VarName, Type)] + +-- | `typecheckExpr` checks that the given `Expr` has the correct types. +typecheckExpr :: MonadError Error m => TypeEnv -> Expr -> m Type +typecheckExpr env = \case + Var x -> case lookup x env of + Nothing -> throwInternalError $ "undefined variable: " ++ unVarName x + Just t -> return t + Lit lit -> return $ literalToType lit + App f e -> do + tf <- typecheckExpr env f + te <- typecheckExpr env e + case tf of + FunTy te' ret | te' == te -> return ret + _ -> throwInternalError $ "wrong type funcall: function = " ++ formatExpr f ++ " and argument = " ++ formatExpr e ++ ", function's type = " ++ formatType tf ++ ", but argument's type = " ++ formatType te + Lam x t e -> FunTy t <$> typecheckExpr ((x, t) : env) e + Let x t e1 e2 -> do + t' <- typecheckExpr env e1 + when (t /= t') $ do + throwInternalError $ "wrong type binding: " ++ formatExpr (Let x t e1 e2) + typecheckExpr ((x, t) : env) e2 + +typecheckToplevelExpr :: MonadError Error m => TypeEnv -> ToplevelExpr -> m Type +typecheckToplevelExpr env = \case + ResultExpr e -> typecheckExpr env e + ToplevelLet x t e cont -> do + t' <- typecheckExpr env e + when (t' /= t) $ do + throwInternalError $ "assigned type is not correct: context = (let " ++ unVarName x ++ ": " ++ formatType t ++ " = " ++ formatExpr e ++ " in ...), expected type = " ++ formatType t ++ ", actual type = " ++ formatType t' + typecheckToplevelExpr ((x, t) : env) cont + ToplevelLetRec f args ret body cont -> do + let t = case args of + [] -> ret + _ -> curryFunTy (map snd args) ret + ret' <- typecheckExpr (reverse args ++ (f, t) : env) body + when (ret' /= ret) $ do + throwInternalError $ "returned type is not correct: context = (let rec " ++ unVarName f ++ " " ++ unwords (map (\(x, t) -> unVarName x ++ ": " ++ formatType t) args) ++ ": " ++ formatType ret ++ " = " ++ formatExpr body ++ " in ...), expected type = " ++ formatType ret ++ ", actual type = " ++ formatType ret' + typecheckToplevelExpr ((f, t) : env) cont + +typecheckProgram :: MonadError Error m => Program -> m Type +typecheckProgram prog = wrapError' "Jikka.Core.Language.TypeCheck.typecheckProgram" $ do + typecheckToplevelExpr [] prog diff --git a/src/Jikka/Core/Language/Util.hs b/src/Jikka/Core/Language/Util.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Language/Util.hs @@ -0,0 +1,354 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +module Jikka.Core.Language.Util where + +import Control.Monad.Identity +import Control.Monad.Writer (execWriter, tell) +import Data.Maybe (isJust) +import Data.Monoid (Dual (..)) +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr + +genType :: MonadAlpha m => m Type +genType = do + i <- nextCounter + return $ VarTy (TypeName ('$' : show i)) + +genVarName :: MonadAlpha m => VarName -> m VarName +genVarName x = do + i <- nextCounter + let base = if unVarName x == "_" then "" else takeWhile (/= '$') (unVarName x) + return $ VarName (base ++ '$' : show i) + +genVarName' :: MonadAlpha m => m VarName +genVarName' = genVarName (VarName "_") + +mapTypeInBuiltin :: (Type -> Type) -> Builtin -> Builtin +mapTypeInBuiltin f = \case + -- arithmetical functions + Negate -> Negate + Plus -> Plus + Minus -> Minus + Mult -> Mult + FloorDiv -> FloorDiv + FloorMod -> FloorMod + CeilDiv -> CeilDiv + CeilMod -> CeilMod + Pow -> Pow + -- advanced arithmetical functions + Abs -> Abs + Gcd -> Gcd + Lcm -> Lcm + Min2 t -> Min2 (f t) + Max2 t -> Max2 (f t) + Iterate t -> Iterate (f t) + -- logical functionslogical + Not -> Not + And -> And + Or -> Or + Implies -> Implies + If t -> If (f t) + -- bitwise functionsbitwise + BitNot -> BitNot + BitAnd -> BitAnd + BitOr -> BitOr + BitXor -> BitXor + BitLeftShift -> BitLeftShift + BitRightShift -> BitRightShift + -- matrix functions + MatAp h w -> MatAp h w + MatZero n -> MatZero n + MatOne n -> MatOne n + MatAdd h w -> MatAdd h w + MatMul h n w -> MatMul h n w + MatPow n -> MatPow n + VecFloorMod n -> VecFloorMod n + MatFloorMod h w -> MatFloorMod h w + -- modular functionsmodular + ModNegate -> ModNegate + ModPlus -> ModPlus + ModMinus -> ModMinus + ModMult -> ModMult + ModInv -> ModInv + ModPow -> ModPow + ModMatAp h w -> ModMatAp h w + ModMatAdd h w -> ModMatAdd h w + ModMatMul h n w -> ModMatMul h n w + ModMatPow n -> ModMatPow n + -- list functionslist + Cons t -> Cons (f t) + Snoc t -> Snoc (f t) + Foldl t1 t2 -> Foldl (f t1) (f t2) + Scanl t1 t2 -> Scanl (f t1) (f t2) + Build t -> Build (f t) + Len t -> Len (f t) + Map t1 t2 -> Map (f t1) (f t2) + Filter t -> Filter (f t) + At t -> At (f t) + SetAt t -> SetAt (f t) + Elem t -> Elem (f t) + Sum -> Sum + Product -> Product + ModSum -> ModSum + ModProduct -> ModProduct + Min1 t -> Min1 (f t) + Max1 t -> Max1 (f t) + ArgMin t -> ArgMin (f t) + ArgMax t -> ArgMax (f t) + All -> All + Any -> Any + Sorted t -> Sorted (f t) + Reversed t -> Reversed (f t) + Range1 -> Range1 + Range2 -> Range2 + Range3 -> Range3 + -- tuple functions + Tuple ts -> Tuple (map f ts) + Proj ts n -> Proj (map f ts) n + -- comparison + LessThan t -> LessThan (f t) + LessEqual t -> LessEqual (f t) + GreaterThan t -> GreaterThan (f t) + GreaterEqual t -> GreaterEqual (f t) + Equal t -> Equal (f t) + NotEqual t -> NotEqual (f t) + -- combinational functions + Fact -> Fact + Choose -> Choose + Permute -> Permute + MultiChoose -> MultiChoose + -- data structures + ConvexHullTrickInit -> ConvexHullTrickInit + ConvexHullTrickInsert -> ConvexHullTrickInsert + ConvexHullTrickGetMin -> ConvexHullTrickGetMin + SegmentTreeInitList semigrp -> SegmentTreeInitList semigrp + SegmentTreeGetRange semigrp -> SegmentTreeGetRange semigrp + SegmentTreeSetPoint semigrp -> SegmentTreeSetPoint semigrp + +-- | `mapExprM'` substitutes exprs using given two functions, which are called in pre-order and post-order. +mapExprM' :: Monad m => ([(VarName, Type)] -> Expr -> m Expr) -> ([(VarName, Type)] -> Expr -> m Expr) -> [(VarName, Type)] -> Expr -> m Expr +mapExprM' pre post env e = do + e <- pre env e + let go = mapExprM' pre post + e <- case e of + Var y -> return $ Var y + Lit lit -> return $ Lit lit + App g e -> App <$> go env g <*> go env e + Lam x t body -> Lam x t <$> go ((x, t) : env) body + Let y t e1 e2 -> Let y t <$> go env e1 <*> go ((y, t) : env) e2 + post env e + +mapExprToplevelExprM' :: Monad m => ([(VarName, Type)] -> Expr -> m Expr) -> ([(VarName, Type)] -> Expr -> m Expr) -> [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr +mapExprToplevelExprM' pre post env = \case + ResultExpr e -> ResultExpr <$> mapExprM' pre post env e + ToplevelLet y t e cont -> + ToplevelLet y t <$> mapExprM' pre post env e <*> mapExprToplevelExprM' pre post ((y, t) : env) cont + ToplevelLetRec g args ret body cont -> + let env' = (g, foldr (FunTy . snd) ret args) : env + in ToplevelLetRec g args ret <$> mapExprM' pre post (reverse args ++ env') body <*> mapExprToplevelExprM' pre post env' cont + +mapExprProgramM' :: Monad m => ([(VarName, Type)] -> Expr -> m Expr) -> ([(VarName, Type)] -> Expr -> m Expr) -> Program -> m Program +mapExprProgramM' pre post = mapExprToplevelExprM' pre post [] + +-- | `mapExprM` is a wrapper of `mapExprM'`. This function works in post-order. +mapExprM :: Monad m => ([(VarName, Type)] -> Expr -> m Expr) -> [(VarName, Type)] -> Expr -> m Expr +mapExprM f = mapExprM' (\_ e -> return e) f + +mapExprToplevelExprM :: Monad m => ([(VarName, Type)] -> Expr -> m Expr) -> [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr +mapExprToplevelExprM f = mapExprToplevelExprM' (\_ e -> return e) f + +mapExprProgramM :: Monad m => ([(VarName, Type)] -> Expr -> m Expr) -> Program -> m Program +mapExprProgramM f = mapExprProgramM' (\_ e -> return e) f + +mapExpr :: ([(VarName, Type)] -> Expr -> Expr) -> [(VarName, Type)] -> Expr -> Expr +mapExpr f env e = runIdentity $ mapExprM (\env e -> return $ f env e) env e + +mapExprToplevelExpr :: ([(VarName, Type)] -> Expr -> Expr) -> [(VarName, Type)] -> ToplevelExpr -> ToplevelExpr +mapExprToplevelExpr f env e = runIdentity $ mapExprToplevelExprM (\env e -> return $ f env e) env e + +mapExprProgram :: ([(VarName, Type)] -> Expr -> Expr) -> Program -> Program +mapExprProgram f prog = runIdentity $ mapExprProgramM (\env e -> return $ f env e) prog + +listSubExprs :: Expr -> [Expr] +listSubExprs e = getDual . execWriter $ mapExprM go [] e + where + go _ e = do + tell $ Dual [e] + return e + +uncurryFunTy :: Type -> ([Type], Type) +uncurryFunTy = \case + (FunTy t t') -> let (ts, ret) = uncurryFunTy t' in (t : ts, ret) + ret -> ([], ret) + +uncurryLam :: Expr -> ([(VarName, Type)], Expr) +uncurryLam = \case + Lam x t body -> let (args, body') = uncurryLam body in ((x, t) : args, body') + body -> ([], body) + +curryApp :: Expr -> (Expr, [Expr]) +curryApp = \case + App f e -> let (f', e') = curryApp f in (f', e' ++ [e]) + f -> (f, []) + +curryFunTy :: [Type] -> Type -> Type +curryFunTy ts ret = foldr FunTy ret ts + +curryLam :: [(VarName, Type)] -> Expr -> Expr +curryLam args body = foldr (uncurry Lam) body args + +uncurryApp :: Expr -> [Expr] -> Expr +uncurryApp = foldl App + +isVectorTy :: Type -> Bool +isVectorTy = isJust . sizeOfVectorTy + +isVectorTy' :: [Type] -> Bool +isVectorTy' = isVectorTy . TupleTy + +sizeOfVectorTy :: Type -> Maybe Int +sizeOfVectorTy = \case + TupleTy ts | all (== IntTy) ts -> Just (length ts) + _ -> Nothing + +isMatrixTy :: Type -> Bool +isMatrixTy = isJust . sizeOfMatrixTy + +isMatrixTy' :: [Type] -> Bool +isMatrixTy' = isMatrixTy . TupleTy + +sizeOfMatrixTy :: Type -> Maybe (Int, Int) +sizeOfMatrixTy = \case + TupleTy ts@(TupleTy ts' : _) | all (== IntTy) ts' && all (== TupleTy ts') ts -> Just (length ts, length ts') + _ -> Nothing + +isConstantTimeBuiltin :: Builtin -> Bool +isConstantTimeBuiltin = \case + -- arithmetical functions + Negate -> True + Plus -> True + Minus -> True + Mult -> True + FloorDiv -> True + FloorMod -> True + CeilDiv -> True + CeilMod -> True + Pow -> True + -- advanced arithmetical functions + Abs -> True + Gcd -> True + Lcm -> True + Min2 _ -> True + Max2 _ -> True + Iterate _ -> False + -- logical functions + Not -> True + And -> True + Or -> True + Implies -> True + If _ -> True + -- bitwise functions + BitNot -> True + BitAnd -> True + BitOr -> True + BitXor -> True + BitLeftShift -> True + BitRightShift -> True + -- matrix functions + MatAp _ _ -> True + MatZero _ -> True + MatOne _ -> True + MatAdd _ _ -> True + MatMul _ _ _ -> True + MatPow _ -> True + VecFloorMod _ -> True + MatFloorMod _ _ -> True + -- modular functions + ModNegate -> True + ModPlus -> True + ModMinus -> True + ModMult -> True + ModInv -> True + ModPow -> True + ModMatAp _ _ -> True + ModMatAdd _ _ -> True + ModMatMul _ _ _ -> True + ModMatPow _ -> True + -- list functions + Cons _ -> False + Snoc _ -> False + Foldl _ _ -> False + Scanl _ _ -> False + Build _ -> False + Len _ -> True + Map _ _ -> False + Filter _ -> False + At _ -> True + SetAt _ -> False + Elem _ -> False + Sum -> False + Product -> False + ModSum -> False + ModProduct -> False + Min1 _ -> False + Max1 _ -> False + ArgMin _ -> False + ArgMax _ -> False + All -> False + Any -> False + Sorted _ -> False + Reversed _ -> False + Range1 -> False + Range2 -> False + Range3 -> False + -- tuple functions + Tuple _ -> True + Proj _ _ -> True + -- comparison + LessThan _ -> True + LessEqual _ -> True + GreaterThan _ -> True + GreaterEqual _ -> True + Equal _ -> True + NotEqual _ -> True + -- combinational functions + Fact -> True + Choose -> True + Permute -> True + MultiChoose -> True + -- data structures + ConvexHullTrickInit -> False + ConvexHullTrickInsert -> False + ConvexHullTrickGetMin -> False + SegmentTreeInitList _ -> False + SegmentTreeGetRange _ -> False + SegmentTreeSetPoint _ -> False + +-- | `isConstantTimeExpr` checks whether given exprs are suitable to propagate. +isConstantTimeExpr :: Expr -> Bool +isConstantTimeExpr = \case + Var _ -> True + Lit _ -> True + e@(App _ _) -> case curryApp e of + (Lit (LitBuiltin f), args) -> isConstantTimeBuiltin f && all isConstantTimeExpr args + _ -> False + Lam _ _ _ -> True + Let _ _ e1 e2 -> isConstantTimeExpr e1 && isConstantTimeExpr e2 + +-- | `replaceLenF` replaces @len(f)@ in an expr with @i + k@. +-- * This assumes that there are no name conflicts. +replaceLenF :: MonadError Error m => VarName -> VarName -> Integer -> Expr -> m Expr +replaceLenF f i k = go + where + go = \case + Len' _ (Var f') | f' == f -> return $ Plus' (Var i) (LitInt' k) + Var y -> return $ Var y + Lit lit -> return $ Lit lit + App g e -> App <$> go g <*> go e + Lam x _ _ | x == i -> throwInternalError "Jikka.Core.Language.Util.replaceLenF: name conflict" + Lam x t body -> Lam x t <$> (if x == f then return body else go body) + Let y _ _ _ | y == i -> throwInternalError "Jikka.Core.Language.Util.replaceLenF: name conflict" + Let y t e1 e2 -> Let y t <$> go e1 <*> (if y == f then return e2 else go e2) diff --git a/src/Jikka/Core/Language/Value.hs b/src/Jikka/Core/Language/Value.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Core/Language/Value.hs @@ -0,0 +1,133 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +module Jikka.Core.Language.Value where + +import Data.Char (toLower) +import Data.List (intercalate) +import qualified Data.Map as M +import Data.Maybe (fromMaybe) +import qualified Data.Vector as V +import Jikka.Common.Error +import Jikka.Common.IOFormat +import Jikka.Common.Matrix +import Jikka.Common.ModInt +import Jikka.Core.Format (formatBuiltinIsolated, formatExpr) +import Jikka.Core.Language.Expr + +data Value + = ValInt Integer + | ValBool Bool + | ValList (V.Vector Value) + | ValTuple [Value] + | ValBuiltin Builtin [Value] + | -- | The `Env` may contain the `ValLambda` cyclicly. + ValLambda (Maybe VarName) Env VarName Type Expr + deriving (Eq, Read) + +type Env = [(VarName, Value)] + +literalToValue :: MonadError Error m => Literal -> m Value +literalToValue = \case + LitBuiltin builtin -> return $ ValBuiltin builtin [] + LitInt n -> return $ ValInt n + LitBool p -> return $ ValBool p + LitNil _ -> return $ ValList V.empty + LitBottom _ err -> throwRuntimeError err + +valueToInt :: MonadError Error m => Value -> m Integer +valueToInt = \case + ValInt n -> return n + val -> throwInternalError $ "not an integer value: " ++ formatValue val + +valueToList :: MonadError Error m => Value -> m (V.Vector Value) +valueToList = \case + ValList xs -> return xs + val -> throwInternalError $ "not a list value: " ++ formatValue val + +valueToIntList :: MonadError Error m => Value -> m [Integer] +valueToIntList xs = mapM valueToInt . V.toList =<< valueToList xs + +valueToBool :: MonadError Error m => Value -> m Bool +valueToBool = \case + ValBool p -> return p + val -> throwInternalError $ "not an boolean value: " ++ formatValue val + +valueToBoolList :: MonadError Error m => Value -> m [Bool] +valueToBoolList xs = mapM valueToBool . V.toList =<< valueToList xs + +valueToTuple :: MonadError Error m => Value -> m [Value] +valueToTuple = \case + ValTuple xs -> return xs + val -> throwInternalError $ "not a tuple value: " ++ formatValue val + +valueToIntPair :: MonadError Error m => Value -> m (Integer, Integer) +valueToIntPair = \case + ValTuple [a, b] -> (,) <$> valueToInt a <*> valueToInt b + val -> throwInternalError $ "not a tuple value: " ++ formatValue val + +valueToVector :: MonadError Error m => Value -> m (V.Vector Integer) +valueToVector = \case + ValTuple x -> V.fromList <$> mapM valueToInt x + val -> throwInternalError $ "not a vector: " ++ formatValue val + +valueToMatrix :: MonadError Error m => Value -> m (Matrix Integer) +valueToMatrix a = do + a <- V.mapM valueToVector . V.fromList =<< valueToTuple a + case makeMatrix a of + Just a -> return a + Nothing -> throwInternalError $ "not a matrix: " ++ show a + +valueFromVector :: V.Vector Integer -> Value +valueFromVector x = ValTuple (map ValInt (V.toList x)) + +valueFromMatrix :: Matrix Integer -> Value +valueFromMatrix f = ValTuple (map (ValTuple . map ValInt . V.toList) (V.toList (unMatrix f))) + +valueToModVector :: MonadError Error m => Integer -> Value -> m (V.Vector ModInt) +valueToModVector m x = V.map (`toModInt` m) <$> valueToVector x + +valueToModMatrix :: MonadError Error m => Integer -> Value -> m (Matrix ModInt) +valueToModMatrix m f = fmap (`toModInt` m) <$> valueToMatrix f + +valueFromModVector :: V.Vector ModInt -> Value +valueFromModVector = valueFromVector . V.map fromModInt + +valueFromModMatrix :: Matrix ModInt -> Value +valueFromModMatrix = valueFromMatrix . fmap fromModInt + +compareValues :: Value -> Value -> Maybe Ordering +compareValues a b = case (a, b) of + (ValInt a, ValInt b) -> Just (compare a b) + (ValBool a, ValBool b) -> Just (compare a b) + (ValList a, ValList b) -> case mconcat <$> zipWithM compareValues (V.toList a) (V.toList b) of + Just EQ -> Just (compare (V.length a) (V.length b)) + ordering -> ordering + (ValTuple a, ValTuple b) -> mconcat <$> zipWithM compareValues a b + (_, _) -> Nothing + +compareValues' :: Value -> Value -> Ordering +compareValues' a b = fromMaybe EQ (compareValues a b) + +minValue :: Value -> Value -> Value +minValue a b = if compareValues' a b == LT then a else b + +maxValue :: Value -> Value -> Value +maxValue a b = if compareValues' a b == GT then a else b + +formatValue :: Value -> String +formatValue = \case + ValInt n -> show n + ValBool p -> map toLower (show p) + ValList xs -> "[" ++ intercalate ", " (map formatValue (V.toList xs)) ++ "]" + ValTuple [x] -> "(" ++ formatValue x ++ ",)" + ValTuple xs -> "(" ++ intercalate ", " (map formatValue xs) ++ ")" + ValBuiltin builtin [] -> formatBuiltinIsolated builtin + ValBuiltin builtin args -> formatBuiltinIsolated builtin ++ "(" ++ intercalate ", " (map formatValue args) ++ ")" + ValLambda _ _ x t body -> formatExpr (Lam x t body) -- Don't show env because it may be cyclic. + +readValueIO :: (MonadError Error m, MonadIO m) => IOFormat -> m ([Value], M.Map String Value) +readValueIO = makeReadValueIO valueToInt ValInt valueToList ValList + +writeValueIO :: (MonadError Error m, MonadIO m) => IOFormat -> M.Map String Value -> Value -> m () +writeValueIO = makeWriteValueIO valueToTuple ValInt valueToInt valueToList diff --git a/src/Jikka/Main.hs b/src/Jikka/Main.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Main.hs @@ -0,0 +1,109 @@ +-- | +-- Module : Jikka.Main +-- Description : is the entry point of the @jikka@ command. / @jikka@ コマンドのエントリポイントです。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Main where + +import Data.Maybe (fromMaybe) +import qualified Data.Text.IO as T +import Data.Version (showVersion) +import Jikka.Common.Error +import Jikka.Common.Format.Error (hPrintError, hPrintErrorWithText) +import qualified Jikka.Main.Subcommand.Convert as Convert +import qualified Jikka.Main.Subcommand.Debug as Debug +import qualified Jikka.Main.Subcommand.Execute as Execute +import Jikka.Main.Target +import Paths_Jikka (version) +import System.Console.GetOpt +import System.Exit (ExitCode (..)) +import System.IO (hPutStr, stderr) + +data Flag + = Help + | Verbose + | Version + | Target String + deriving (Eq, Ord, Show, Read) + +data Options = Options + { verbose :: Bool, + target :: Maybe Target + } + deriving (Eq, Ord, Show, Read) + +defaultOptions :: Options +defaultOptions = + Options + { verbose = False, + target = Nothing + } + +header :: String -> String +header progName = "Usage: " ++ progName ++ " [convert | debug | execute] FILE" + +options :: [OptDescr Flag] +options = + [ Option ['h', '?'] ["help"] (NoArg Help) "", + Option ['v'] ["verbose"] (NoArg Verbose) "", + Option [] ["version"] (NoArg Version) "", + Option [] ["target"] (ReqArg Target "TARGET") "\"python\", \"rpython\", \"core\" or \"cxx\"" + ] + +main :: String -> [String] -> IO ExitCode +main name args = do + let usage = usageInfo (header name) options + case getOpt Permute options args of + (parsed, _, []) | Help `elem` parsed -> do + putStr usage + return ExitSuccess + (parsed, _, []) | Version `elem` parsed -> do + putStrLn $ showVersion version + return ExitSuccess + (parsed, [subcmd, path], []) -> case parseFlags name parsed of + Left err -> do + hPrintError stderr err + return $ ExitFailure 1 + Right opts -> do + result <- runExceptT $ runSubcommand subcmd opts path + case result of + Left err -> do + text <- liftIO $ T.readFile path + hPrintErrorWithText stderr text err + return $ ExitFailure 1 + Right () -> do + return ExitSuccess + (_, _, errors) | errors /= [] -> do + forM_ errors $ \msg -> do + let err = WithGroup CommandLineError (Error msg) + hPrintError stderr err + return $ ExitFailure 1 + _ -> do + hPutStr stderr usage + return $ ExitFailure 1 + +parseFlags :: String -> [Flag] -> Either Error Options +parseFlags _ = go defaultOptions + where + go :: Options -> [Flag] -> Either Error Options + go opts [] = Right opts + go opts (flag : flags) = case flag of + Help -> throwCommandLineError "parseFlags is not called when --help is specified" + Version -> throwCommandLineError "parseFlags is not called when --version is specified" + Verbose -> go (opts {verbose = True}) flags + Target target -> do + target <- parseTarget target + go (opts {target = Just target}) flags + +runSubcommand :: String -> Options -> FilePath -> ExceptT Error IO () +runSubcommand subcmd opts path = case subcmd of + "convert" -> do + input <- liftIO $ T.readFile path + output <- liftEither $ Convert.run (fromMaybe CPlusPlusTarget (target opts)) path input + liftIO $ T.putStr output + "debug" -> Debug.run path + "execute" -> Execute.run (fromMaybe CoreTarget (target opts)) path + _ -> throwCommandLineError $ "undefined subcommand: " ++ show subcmd diff --git a/src/Jikka/Main/Subcommand/Convert.hs b/src/Jikka/Main/Subcommand/Convert.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Main/Subcommand/Convert.hs @@ -0,0 +1,61 @@ +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Subcommand.Convert +-- Description : is the entry point of @convert@ subcommand. / @convert@ サブコマンドのエントリポイントです。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Main.Subcommand.Convert (run) where + +import Data.Text (Text, pack) +import qualified Jikka.CPlusPlus.Convert as FromCore +import qualified Jikka.CPlusPlus.Format as FormatCPlusPlus +import Jikka.Common.Alpha +import Jikka.Common.Error +import qualified Jikka.Core.Convert as Convert +import qualified Jikka.Core.Format as FormatCore +import Jikka.Main.Target +import qualified Jikka.Python.Convert.ToRestrictedPython as ToRestrictedPython +import qualified Jikka.Python.Parse as ParsePython +import qualified Jikka.RestrictedPython.Convert as ToCore +import qualified Jikka.RestrictedPython.Format as FormatRestrictedPython + +runPython :: FilePath -> Text -> Either Error Text +runPython path input = flip evalAlphaT 0 $ do + prog <- ParsePython.run path input + return . pack $ show prog -- TODO + +runRestrictedPython :: FilePath -> Text -> Either Error Text +runRestrictedPython path input = flip evalAlphaT 0 $ do + prog <- ParsePython.run path input + prog <- ToRestrictedPython.run prog + (prog, _) <- ToCore.run' prog + FormatRestrictedPython.run prog + +runCore :: FilePath -> Text -> Either Error Text +runCore path input = flip evalAlphaT 0 $ do + prog <- ParsePython.run path input + prog <- ToRestrictedPython.run prog + (prog, _) <- ToCore.run prog + prog <- Convert.run prog + FormatCore.run prog + +runCPlusPlus :: FilePath -> Text -> Either Error Text +runCPlusPlus path input = flip evalAlphaT 0 $ do + prog <- ParsePython.run path input + prog <- ToRestrictedPython.run prog + (prog, format) <- ToCore.run prog + prog <- Convert.run prog + resetAlphaT 0 -- to make generated C++ code cleaner + prog <- FromCore.run prog format + FormatCPlusPlus.run prog + +run :: Target -> FilePath -> Text -> Either Error Text +run = \case + PythonTarget -> runPython + RestrictedPythonTarget -> runRestrictedPython + CoreTarget -> runCore + CPlusPlusTarget -> runCPlusPlus diff --git a/src/Jikka/Main/Subcommand/Debug.hs b/src/Jikka/Main/Subcommand/Debug.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Main/Subcommand/Debug.hs @@ -0,0 +1,20 @@ +module Jikka.Main.Subcommand.Debug (run) where + +import qualified Data.Text.IO as T (putStrLn, readFile) +import Jikka.Common.Alpha +import Jikka.Common.Error +import qualified Jikka.Common.IOFormat as IOFormat +import qualified Jikka.Core.Format as FormatCore +import qualified Jikka.Python.Convert.ToRestrictedPython as ToRestrictedPython +import qualified Jikka.Python.Parse as ParsePython +import qualified Jikka.RestrictedPython.Convert as ToCore + +run :: FilePath -> ExceptT Error IO () +run path = flip evalAlphaT 0 $ do + prog <- liftIO $ T.readFile path + prog <- ParsePython.run path prog + prog <- ToRestrictedPython.run prog + (prog, format) <- ToCore.run prog + prog <- FormatCore.run prog + liftIO $ T.putStrLn prog + liftIO $ putStr (IOFormat.formatIOFormat format) diff --git a/src/Jikka/Main/Subcommand/Execute.hs b/src/Jikka/Main/Subcommand/Execute.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Main/Subcommand/Execute.hs @@ -0,0 +1,59 @@ +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Subcommand.Execute +-- Description : is the entry point of @execute@ subcommand. / @execute@ サブコマンドのエントリポイントです。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Main.Subcommand.Execute (run) where + +import Control.Monad.Except +import qualified Data.Text.IO as T (readFile) +import Jikka.Common.Alpha +import Jikka.Common.Error +import qualified Jikka.Core.Convert as ConvertCore +import qualified Jikka.Core.Evaluate as EvaluateCore +import qualified Jikka.Core.Language.Value as ValueCore +import Jikka.Main.Target +import qualified Jikka.Python.Convert.ToRestrictedPython as ToRestrictedPython +import qualified Jikka.Python.Parse as FromPython +import qualified Jikka.RestrictedPython.Convert as ToCore +import qualified Jikka.RestrictedPython.Evaluate as EvaluateRestrictedPython +import qualified Jikka.RestrictedPython.Language.Value as ValueRestrictedPythong + +runPython :: FilePath -> ExceptT Error IO () +runPython _ = throwCommandLineError "cannot execute Python" + +runRestrictedPython :: FilePath -> ExceptT Error IO () +runRestrictedPython path = flip evalAlphaT 0 $ do + prog <- liftIO $ T.readFile path + prog <- liftEither $ FromPython.run path prog + prog <- ToRestrictedPython.run prog + (prog, format) <- ToCore.run' prog + (args, env) <- ValueRestrictedPythong.readValueIO format + result <- EvaluateRestrictedPython.run prog args + ValueRestrictedPythong.writeValueIO format env result + +runCore :: FilePath -> ExceptT Error IO () +runCore path = flip evalAlphaT 0 $ do + prog <- liftIO $ T.readFile path + prog <- liftEither $ FromPython.run path prog + prog <- ToRestrictedPython.run prog + (prog, format) <- ToCore.run prog + prog <- ConvertCore.run prog + (args, env) <- ValueCore.readValueIO format + result <- EvaluateCore.run prog args + ValueCore.writeValueIO format env result + +runCPlusPlus :: FilePath -> ExceptT Error IO () +runCPlusPlus _ = throwCommandLineError "cannot execute C++" + +run :: Target -> FilePath -> ExceptT Error IO () +run = \case + PythonTarget -> runPython + RestrictedPythonTarget -> runRestrictedPython + CoreTarget -> runCore + CPlusPlusTarget -> runCPlusPlus diff --git a/src/Jikka/Main/Target.hs b/src/Jikka/Main/Target.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Main/Target.hs @@ -0,0 +1,20 @@ +{-# LANGUAGE LambdaCase #-} + +module Jikka.Main.Target where + +import Jikka.Common.Error + +data Target + = PythonTarget + | RestrictedPythonTarget + | CoreTarget + | CPlusPlusTarget + deriving (Eq, Ord, Show, Read) + +parseTarget :: String -> Either Error Target +parseTarget = \case + "python" -> return PythonTarget + "rpython" -> return RestrictedPythonTarget + "core" -> return CoreTarget + "cxx" -> return CPlusPlusTarget + s -> throwCommandLineError $ "invalid target: " ++ s diff --git a/src/Jikka/Python/Convert/ToRestrictedPython.hs b/src/Jikka/Python/Convert/ToRestrictedPython.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Python/Convert/ToRestrictedPython.hs @@ -0,0 +1,274 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.Python.Convert.ToRestrictedPython +-- Description : converts AST of the standard Python to AST of our restricted Python. / 標準の Python の抽象構文木を我々の restricted Python の抽象構文木に変換します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Python.Convert.ToRestrictedPython + ( run, + ) +where + +import Control.Monad.Except +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Common.Location +import qualified Jikka.Python.Language.Expr as X +import qualified Jikka.RestrictedPython.Language.Expr as Y +import qualified Jikka.RestrictedPython.Language.Util as Y (genType) + +-- --------------------------------------------------------------------------- +-- convert AST + +runIdent :: X.Ident' -> Y.VarName' +runIdent (WithLoc loc (X.Ident x)) = WithLoc' (Just loc) (Y.VarName x) + +runAttribute :: X.Ident' -> Y.Attribute' +runAttribute (WithLoc loc (X.Ident x)) = WithLoc' (Just loc) (Y.UnresolvedAttribute (Y.AttributeName x)) + +runType :: (MonadAlpha m, MonadError Error m) => X.Type' -> m Y.Type +runType t = wrapAt (loc t) $ case value t of + X.Constant (X.ConstString _) -> Y.genType + X.Constant X.ConstNone -> return Y.NoneTy + X.Name (WithLoc _ (X.Ident "int")) -> return Y.IntTy + X.Name (WithLoc _ (X.Ident "bool")) -> return Y.BoolTy + X.Subscript (WithLoc _ (X.Name (WithLoc _ (X.Ident f)))) e -> case (f, e) of + ("List", _) -> Y.ListTy <$> runType e + ("Iterator", _) -> Y.ListTy <$> runType e + ("Sequence", _) -> Y.ListTy <$> runType e + ("Tuple", WithLoc _ (X.Tuple es)) -> Y.TupleTy <$> mapM runType es + ("Tuple", _) -> Y.TupleTy . (: []) <$> runType e + ("Callable", WithLoc _ (X.Tuple [WithLoc _ (X.List es), e])) -> do + ts <- mapM runType es + t <- runType e + return $ Y.CallableTy ts t + _ -> throwSemanticError ("not a type: " ++ show t) + _ -> throwSemanticError ("not a type: " ++ show t) + +runMaybeType :: (MonadAlpha m, MonadError Error m) => Maybe X.Type' -> m Y.Type +runMaybeType Nothing = Y.genType +runMaybeType (Just t) = runType t + +runConstant :: MonadError Error m => X.Constant -> m Y.Constant +runConstant = \case + X.ConstNone -> return Y.ConstNone + X.ConstInt n -> return $ Y.ConstInt n + X.ConstBool p -> return $ Y.ConstBool p + e -> throwSemanticError ("unsupported constant: " ++ show e) + +runTargetName :: (MonadAlpha m, MonadError Error m) => X.Expr' -> m Y.VarName' +runTargetName e = case value e of + X.Name x -> return $ runIdent x + _ -> throwSemanticErrorAt (loc e) ("not an assignment target: " ++ show e) + +runTarget :: (MonadAlpha m, MonadError Error m) => X.Expr' -> m Y.Target' +runTarget e = + WithLoc' (Just (loc e)) <$> case value e of + X.Subscript f index -> Y.SubscriptTrg <$> runTarget f <*> runExpr index + X.Name _ -> Y.NameTrg <$> runTargetName e + X.Tuple es -> Y.TupleTrg <$> mapM runTarget es + _ -> throwSemanticErrorAt (loc e) ("not an assignment target: " ++ show e) + +runTargetIdent :: MonadError Error m => X.Expr' -> m Y.VarName' +runTargetIdent e = case value e of + X.Name x -> return $ runIdent x + _ -> throwSemanticErrorAt (loc e) ("not an simple assignment target: " ++ show e) + +runComprehension :: (MonadAlpha m, MonadError Error m) => [X.Comprehension] -> m Y.Comprehension +runComprehension = \case + [comp] -> do + x <- runTarget (X.compTarget comp) + iter <- runExpr (X.compIter comp) + ifs <- mapM runExpr (X.compIfs comp) + return $ Y.Comprehension x iter ifs + comp -> throwSemanticError ("many comprehensions are unsupported: " ++ show comp) + +runArguments :: (MonadAlpha m, MonadError Error m) => X.Arguments -> m [(Y.VarName', Y.Type)] +runArguments = \case + X.Arguments + { X.argsPosonlyargs = [], + X.argsArgs = args, + X.argsVarargs = Nothing, + X.argsKwonlyargs = [], + X.argsKwDefaults = [], + X.argsKwarg = Nothing, + X.argsDefaults = [] + } -> do + forM args $ \(x, t) -> do + let x' = runIdent x + t <- runMaybeType t + return (x', t) + args -> throwSemanticError ("unsupported arguments: " ++ show args) + +runCompareExpr :: (MonadAlpha m, MonadError Error m) => X.Expr' -> [(X.CmpOp, X.Expr')] -> m Y.Expr +runCompareExpr e1 ops = value' <$> (runExpr e1 >>= (`go` ops)) + where + withLoc = WithLoc' (Just (loc e1)) + go :: (MonadAlpha m, MonadError Error m) => Y.Expr' -> [(X.CmpOp, X.Expr')] -> m Y.Expr' + go e1 = \case + [] -> return . withLoc $ Y.Constant (Y.ConstBool True) + [(op, e2)] -> withLoc <$> (Y.Compare e1 <$> (Y.CmpOp' op <$> Y.genType) <*> runExpr e2) + (op, e2) : ops -> do + t <- Y.genType + e2 <- runExpr e2 + cont <- go e2 ops + return . withLoc $ Y.BoolOp (withLoc (Y.Compare e1 (Y.CmpOp' op t) e2)) Y.And cont + +runExpr :: (MonadAlpha m, MonadError Error m) => X.Expr' -> m Y.Expr' +runExpr e = + WithLoc' (Just (loc e)) <$> case value e of + X.BoolOp e1 op e2 -> Y.BoolOp <$> runExpr e1 <*> return op <*> runExpr e2 + X.BinOp e1 op e2 -> Y.BinOp <$> runExpr e1 <*> return op <*> runExpr e2 + X.UnaryOp op e -> Y.UnaryOp op <$> runExpr e + X.Lambda args body -> Y.Lambda <$> runArguments args <*> runExpr body + X.IfExp e1 e2 e3 -> Y.IfExp <$> runExpr e1 <*> runExpr e2 <*> runExpr e3 + X.ListComp e comp -> Y.ListComp <$> runExpr e <*> runComprehension comp + X.GeneratorExp e comp -> Y.ListComp <$> runExpr e <*> runComprehension comp + X.Compare e1 e2 -> runCompareExpr e1 e2 + X.Call f args [] -> Y.Call <$> runExpr f <*> mapM runExpr args + X.Constant const -> Y.Constant <$> runConstant const + X.Attribute e x -> Y.Attribute <$> runExpr e <*> pure (runAttribute x) + X.Subscript e1 e2 -> case value e2 of + X.Slice from to step -> Y.SubscriptSlice <$> runExpr e1 <*> mapM runExpr from <*> mapM runExpr to <*> mapM runExpr step + _ -> Y.Subscript <$> runExpr e1 <*> runExpr e2 + X.Starred e -> Y.Starred <$> runExpr e + X.Name x -> return $ Y.Name (runIdent x) + X.List es -> Y.List <$> Y.genType <*> mapM runExpr es + X.Tuple es -> Y.Tuple <$> mapM runExpr es + _ -> throwSemanticErrorAt (loc e) ("unsupported expr: " ++ show e) + +runStatement :: (MonadAlpha m, MonadError Error m) => X.Statement' -> m [Y.Statement] +runStatement stmt = wrapAt (loc stmt) $ case value stmt of + X.FunctionDef _ _ _ _ _ -> throwSemanticError "def statement is not allowed in def statement" + X.AsyncFunctionDef _ _ _ _ _ -> throwSemanticError "async-def statement is not allowed in def statement" + X.ClassDef _ _ _ _ _ -> throwSemanticError "class statement is not allowed in def statement" + X.Return e -> do + e <- case e of + Nothing -> return . WithLoc' (Just (loc stmt)) $ Y.Constant Y.ConstNone + Just e -> runExpr e + return [Y.Return e] + X.Delete _ -> throwSemanticErrorAt (loc stmt) "del statement is not allowed in def statement" + X.Assign xs e -> case xs of + [] -> return [] + [x] -> do + x <- runTarget x + t <- Y.genType + e <- runExpr e + return [Y.AnnAssign x t e] + _ -> throwSemanticError "assign statement with multiple targets is not allowed in def statement" + X.AugAssign x op e -> do + x <- runTarget x + e <- runExpr e + return [Y.AugAssign x op e] + X.AnnAssign x t e -> case e of + Nothing -> throwSemanticError "annotated assignment statement without value is not allowed in def statement" + Just e -> do + x <- runTargetIdent x + t <- runType t + e <- runExpr e + return [Y.AnnAssign (WithLoc' (loc' x) (Y.NameTrg x)) t e] + X.For x e body orelse -> do + x <- runTarget x + e <- runExpr e + body <- runStatements body + orelse <- runStatements orelse + return $ Y.For x e body : orelse + X.AsyncFor _ _ _ _ -> throwSemanticError "async-for statement is not allowed in def statement" + X.While _ _ _ -> throwSemanticError "while statement is not allowed in def statement" + X.If e body1 body2 -> do + e <- runExpr e + body1 <- runStatements body1 + body2 <- runStatements body2 + return [Y.If e body1 body2] + X.With _ _ -> throwSemanticError "with statement is not allowed in def statement" + X.AsyncWith _ _ -> throwSemanticError "async-with statement is not allowed in def statement" + X.Raise _ _ -> throwSemanticError "raise statement is not allowed in def statement" + X.Try _ _ _ _ -> throwSemanticError "try statement is not allowed in def statement" + X.Assert e _ -> do + e <- runExpr e + return [Y.Assert e] + X.Import _ -> throwSemanticError "import statement is not allowed in def statement" + X.ImportFrom _ _ -> throwSemanticError "import-from statement is not allowed in def statement" + X.Global _ -> throwSemanticError "global statement is not allowed in def statement" + X.Nonlocal _ -> throwSemanticError "nonlocal statement is not allowed in def statement" + X.Expr' e -> do + e <- runExpr e + return [Y.Expr' e] + X.Pass -> return [] + X.Break -> throwSemanticError "break statement is not allowed in def statement" + X.Continue -> throwSemanticError "continue statement is not allowed in def statement" + +runStatements :: (MonadAlpha m, MonadError Error m) => [X.Statement'] -> m [Y.Statement] +runStatements stmts = do + stmts <- mapM (catchError' . runStatement) stmts + concat <$> reportErrors stmts + +runToplevelStatement :: (MonadAlpha m, MonadError Error m) => X.Statement' -> m [Y.ToplevelStatement] +runToplevelStatement stmt = wrapAt (loc stmt) $ case value stmt of + X.FunctionDef f args body decorators ret -> case decorators of + [] -> do + let f' = runIdent f + args <- runArguments args + body <- runStatements body + ret <- runMaybeType ret + return [Y.ToplevelFunctionDef f' args ret body] + _ -> throwSemanticError "def statement with decorators is not allowed at toplevel" + X.AsyncFunctionDef _ _ _ _ _ -> throwSemanticError "async-def statement is not allowed at toplevel" + X.ClassDef _ _ _ _ _ -> throwSemanticError "class statement is not allowed at toplevel" + X.Return _ -> throwSemanticError "retrun statement is not allowed at toplevel" + X.Delete _ -> throwSemanticError "del statement is not allowed at toplevel" + X.Assign xs e -> case xs of + [] -> return [] + [x] -> do + x <- runTargetIdent x + t <- Y.genType + e <- runExpr e + return [Y.ToplevelAnnAssign x t e] + _ -> throwSemanticError "assignment statement with multiple targets is not allowed at toplevel" + X.AugAssign _ _ _ -> throwSemanticError "augumented assignment statement is not allowed at toplevel" + X.AnnAssign x t e -> case e of + Nothing -> throwSemanticError "annotated assignment statement without value is not allowed at toplevel" + Just e -> do + x <- runTargetIdent x + t <- runType t + e <- runExpr e + return [Y.ToplevelAnnAssign x t e] + X.For _ _ _ _ -> throwSemanticError "for statement is not allowed at toplevel" + X.AsyncFor _ _ _ _ -> throwSemanticError "async-for statement is not allowed at toplevel" + X.While _ _ _ -> throwSemanticError "while statement is not allowed at toplevel" + X.If e body1 body2 -> case (e, body1, body2) of + ( WithLoc _ (X.Compare (WithLoc _ (X.Name (WithLoc _ (X.Ident "__name__")))) [(X.Eq', WithLoc _ (X.Constant (X.ConstString "__main__")))]), + [WithLoc _ (X.Expr' (WithLoc _ (X.Call (WithLoc _ (X.Name (WithLoc _ (X.Ident "main")))) [] [])))], + [] + ) -> return [] + _ -> throwSemanticError "only `if __name__ == \"__main__\": main()' is allowed for if statements at toplevel" + X.With _ _ -> throwSemanticError "with statement is not allowed at toplevel" + X.AsyncWith _ _ -> throwSemanticError "async-with statement is not allowed at toplevel" + X.Raise _ _ -> throwSemanticError "raise statement is not allowed at toplevel" + X.Try _ _ _ _ -> throwSemanticError "try statement is not allowed at toplevel" + X.Assert e _ -> do + e <- runExpr e + return [Y.ToplevelAssert e] + X.Import _ -> return [] + X.ImportFrom _ _ -> return [] + X.Global _ -> throwSemanticError "global statement is not allowed at toplevel" + X.Nonlocal _ -> throwSemanticError "nonlocal statement is not allowed at toplevel" + X.Expr' e -> case e of + WithLoc _ (X.Call (WithLoc _ (X.Name (WithLoc _ (X.Ident "main")))) [] []) -> return [] + _ -> throwSemanticError "only `main()' is allowed for expression statements at toplevel" + X.Pass -> return [] + X.Break -> throwSemanticError "break statement is not allowed at toplevel" + X.Continue -> throwSemanticError "continue statement is not allowed at toplevel" + +runProgram :: (MonadAlpha m, MonadError Error m) => X.Program -> m Y.Program +runProgram stmts = do + stmts <- mapM (catchError' . runToplevelStatement) stmts + concat <$> reportErrors stmts + +run :: (MonadAlpha m, MonadError Error m) => X.Program -> m Y.Program +run prog = wrapError' "Failed at Jikka.Python.Convert.ToplevelDecl" $ runProgram prog diff --git a/src/Jikka/Python/Language/Expr.hs b/src/Jikka/Python/Language/Expr.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Python/Language/Expr.hs @@ -0,0 +1,236 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} + +-- | +-- Module : Jikka.Python.Language.Expr +-- Description : contains data types of the standard Python. / 標準の Python のためのデータ型を含みます。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- `Jikka.Python.Language.Expr` module has the basic data types for the standard Python. +-- See the Python's @ast@ module () for reference. +module Jikka.Python.Language.Expr where + +import Data.Int (Int8) +import Data.String (IsString) +import Jikka.Common.Location + +newtype Ident = Ident String deriving (Eq, Ord, Show, Read, IsString) + +unIdent :: Ident -> String +unIdent (Ident x) = x + +type Ident' = WithLoc Ident + +data Constant + = ConstNone + | ConstInt Integer + | ConstBool Bool + | ConstString String + | ConstBytes [Int8] + | ConstFloat Double + | ConstImaginary Double + deriving (Eq, Ord, Show, Read) + +data Statement + = FunctionDef Ident' Arguments [Statement'] [Decorator] (Maybe Type') + | AsyncFunctionDef Ident' Arguments [Statement'] [Decorator] (Maybe Type') + | ClassDef Ident' [Expr'] [Keyword'] [Statement'] [Decorator] + | Return (Maybe Expr') + | Delete [Target'] + | Assign [Target'] Expr' + | AugAssign Target' Operator Expr' + | AnnAssign Target' Type' (Maybe Expr') + | -- | @For target iter body orelse@ + For Target' Expr' [Statement'] [Statement'] + | AsyncFor Target' Expr' [Statement'] [Statement'] + | -- | @While test body orelse@ + While Expr' [Statement'] [Statement'] + | If Expr' [Statement'] [Statement'] + | With [WithItem] [Statement'] + | AsyncWith [WithItem] [Statement'] + | -- | @Raise exc cause@ represents @raise exc from cause@. + Raise (Maybe Expr') (Maybe Expr') + | -- | @Try body handlers orelse finalbody@ + Try [Statement'] [ExceptHandler'] [Statement'] [Statement'] + | -- | @Assert test msg@ + Assert Expr' (Maybe Expr') + | Import [Alias] + | ImportFrom [Ident'] [Alias] + | Global [Ident'] + | Nonlocal [Ident'] + | Expr' Expr' + | Pass + | Break + | Continue + deriving (Eq, Ord, Show, Read) + +type Statement' = WithLoc Statement + +data Expr + = BoolOp Expr' BoolOp Expr' + | -- | produced by the walrus operator @:=@ + NamedExpr Target' Expr' + | BinOp Expr' Operator Expr' + | UnaryOp UnaryOp Expr' + | Lambda Arguments Expr' + | -- | @IfExp test body orelse@ + IfExp Expr' Expr' Expr' + | -- | NULL key is for @**d@. + Dict [(Maybe Expr', Expr')] + | Set [Expr'] + | ListComp Expr' [Comprehension] + | SetComp Expr' [Comprehension] + | DictComp Expr' Expr' [Comprehension] + | GeneratorExp Expr' [Comprehension] + | Await Expr' + | Yield (Maybe Expr') + | YieldFrom Expr' + | Compare Expr' [(CmpOp, Expr')] + | Call Expr' [Expr'] [Keyword'] + | -- | @FormattedValue value conversion format_spec@ for f-strings + FormattedValue Expr' (Maybe Char) (Maybe Expr') + | JoinedStr [Expr'] -- for f-strings + | Constant Constant + | -- | can appear in assignment context + Attribute Expr' Ident' + | -- | can appear in assignment context + Subscript Expr' Expr' + | -- | can appear in assignment context + Starred Expr' + | -- | can appear in assignment context + Name Ident' + | -- | can appear in assignment context + List [Expr'] + | -- | can appear in assignment context + Tuple [Expr'] + | -- | @Slice lower upper step@ can appear only in Subscript + Slice (Maybe Expr') (Maybe Expr') (Maybe Expr') + deriving (Eq, Ord, Show, Read) + +type Expr' = WithLoc Expr + +type Target = Expr' + +type Target' = Expr' + +type Type = Expr' + +type Type' = Expr' + +type Decorator = Expr' + +type Decorator' = Expr' + +data ExprContext = Load | Store | Del + deriving (Eq, Ord, Show, Read) + +data BoolOp + = And + | Or + | -- | our extension + Implies + deriving (Eq, Ord, Show, Read) + +data Operator + = Add + | Sub + | Mult + | MatMult + | Div + | FloorDiv + | FloorMod + | -- | our extension + CeilDiv + | -- | our extension + CeilMod + | Pow + | BitLShift + | BitRShift + | BitOr + | BitXor + | BitAnd + | -- | our extension + Max + | -- | our extension + Min + deriving (Eq, Ord, Show, Read) + +data UnaryOp + = -- | on int + Invert + | -- | on bool + Not + | UAdd + | USub + deriving (Eq, Ord, Show, Read) + +data CmpOp + = Eq' + | NotEq + | Lt + | LtE + | Gt + | GtE + | Is + | IsNot + | In + | NotIn + deriving (Eq, Ord, Show, Read) + +-- | @Comprehension target iter ifs is_async@ +data Comprehension = Comprehension + { compTarget :: Target', + compIter :: Expr', + compIfs :: Maybe Expr' + } + deriving (Eq, Ord, Show, Read) + +data ExceptHandler = ExceptHandler + { exchType :: Maybe Type', + exchName :: Maybe Ident', + exchBody :: [Statement'] + } + deriving (Eq, Ord, Show, Read) + +type ExceptHandler' = WithLoc ExceptHandler + +data Arguments = Arguments + { argsPosonlyargs :: [Arg], + argsArgs :: [Arg], + argsVarargs :: Maybe Arg, + argsKwonlyargs :: [Arg], + argsKwDefaults :: [Expr'], + argsKwarg :: Maybe Arg, + argsDefaults :: [Expr'] + } + deriving (Eq, Ord, Show, Read) + +emptyArguments :: Arguments +emptyArguments = + Arguments + { argsPosonlyargs = [], + argsArgs = [], + argsVarargs = Nothing, + argsKwonlyargs = [], + argsKwDefaults = [], + argsKwarg = Nothing, + argsDefaults = [] + } + +type Arg = (Ident', Maybe Type') + +-- | NULL identifier is for @**kwargs@. +type Keyword = (Maybe Ident', Expr') + +type Keyword' = WithLoc Keyword + +-- | @(name, asname)@. `Alias` is used for `Import` and `ImportFrom`. +type Alias = (Ident', Maybe Ident') + +-- | @(context_expr, optional_vars)@ +type WithItem = (Expr', Maybe Expr') + +type Program = [Statement'] diff --git a/src/Jikka/Python/Language/Util.hs b/src/Jikka/Python/Language/Util.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Python/Language/Util.hs @@ -0,0 +1,110 @@ +module Jikka.Python.Language.Util where + +import Control.Monad.Identity +import Jikka.Common.Location +import Jikka.Python.Language.Expr + +constIntExp :: Integer -> Expr +constIntExp = Constant . ConstInt + +constBoolExp :: Bool -> Expr +constBoolExp = Constant . ConstBool + +mapExprArgumentsM :: Monad m => (Expr' -> m Expr') -> Arguments -> m Arguments +mapExprArgumentsM f args = do + kwDefaults <- mapM (mapExprM f) (argsKwDefaults args) + defaults <- mapM (mapExprM f) (argsDefaults args) + return $ + args + { argsKwDefaults = kwDefaults, + argsDefaults = defaults + } + +mapExprComprehensionM :: Monad m => (Expr' -> m Expr') -> Comprehension -> m Comprehension +mapExprComprehensionM f comp = do + iter <- mapExprM f (compIter comp) + target <- mapExprM f (compIter comp) + ifs <- mapM (mapExprM f) (compIfs comp) + return $ Comprehension {compTarget = target, compIter = iter, compIfs = ifs} + +mapExprComprehensionsM :: Monad m => (Expr' -> m Expr') -> [Comprehension] -> m [Comprehension] +mapExprComprehensionsM f = mapM (mapExprComprehensionM f) + +mapExprKeywordsM :: Monad m => (Expr' -> m Expr') -> [Keyword'] -> m [Keyword'] +mapExprKeywordsM f kwargs = mapM (\(WithLoc loc (k, v)) -> WithLoc loc . (,) k <$> mapExprM f v) kwargs + +mapExprM :: Monad m => (Expr' -> m Expr') -> Expr' -> m Expr' +mapExprM f = go + where + go e0 = + f . WithLoc (loc e0) =<< case value e0 of + BoolOp e1 op e2 -> BoolOp <$> go e1 <*> pure op <*> go e2 + NamedExpr x e -> NamedExpr <$> go x <*> go e + BinOp e1 op e2 -> BinOp <$> go e1 <*> pure op <*> go e2 + UnaryOp op e -> UnaryOp op <$> go e + Lambda args body -> Lambda <$> mapExprArgumentsM f args <*> mapExprM f body + IfExp e1 e2 e3 -> IfExp <$> go e1 <*> go e2 <*> go e3 + Dict es -> Dict <$> mapM (\(k, v) -> (,) <$> traverse go k <*> go v) es + Set es -> Set <$> mapM go es + ListComp e comps -> ListComp <$> mapExprM f e <*> mapExprComprehensionsM f comps + SetComp e comps -> SetComp <$> mapExprM f e <*> mapExprComprehensionsM f comps + DictComp k v comps -> DictComp <$> mapExprM f k <*> mapExprM f v <*> mapExprComprehensionsM f comps + GeneratorExp e comps -> GeneratorExp <$> mapExprM f e <*> mapExprComprehensionsM f comps + Await e -> Await <$> go e + Yield e -> Yield <$> traverse go e + YieldFrom e -> YieldFrom <$> go e + Compare e es -> Compare <$> go e <*> mapM (\(op, e) -> (,) op <$> go e) es + Call g args kwargs -> Call <$> go g <*> mapM go args <*> mapExprKeywordsM f kwargs + FormattedValue e1 c e2 -> FormattedValue <$> go e1 <*> pure c <*> traverse go e2 + JoinedStr es -> JoinedStr <$> mapM go es + Constant constant -> return $ Constant constant + Attribute e x -> Attribute <$> go e <*> pure x + Subscript e1 e2 -> Subscript <$> go e1 <*> go e2 + Starred e -> Starred <$> go e + Name x -> return $ Name x + List es -> List <$> mapM go es + Tuple es -> Tuple <$> mapM go es + Slice e1 e2 e3 -> Slice <$> traverse go e1 <*> traverse go e2 <*> traverse go e3 + +mapExprExceptHanderM :: Monad m => (Expr' -> m Expr') -> ExceptHandler' -> m ExceptHandler' +mapExprExceptHanderM f (WithLoc loc handler) = do + body <- mapExprStatementsM f (exchBody handler) + return $ WithLoc loc (handler {exchBody = body}) + +mapExprStatementM :: Monad m => (Expr' -> m Expr') -> Statement' -> m Statement' +mapExprStatementM f stmt = + WithLoc (loc stmt) <$> case value stmt of + FunctionDef g args body decorators ret -> FunctionDef g args <$> mapExprStatementsM f body <*> mapM (mapExprM f) decorators <*> pure ret + AsyncFunctionDef g args body decorators ret -> AsyncFunctionDef g args <$> mapExprStatementsM f body <*> mapM (mapExprM f) decorators <*> pure ret + ClassDef name bases keywords body decorators -> ClassDef name <$> mapM (mapExprM f) bases <*> mapExprKeywordsM f keywords <*> mapExprStatementsM f body <*> mapM (mapExprM f) decorators + Return e -> Return <$> traverse (mapExprM f) e + Delete xs -> Delete <$> mapM (mapExprM f) xs + Assign xs e -> Assign <$> mapM (mapExprM f) xs <*> mapExprM f e + AugAssign x op e -> AugAssign <$> mapExprM f x <*> pure op <*> mapExprM f e + AnnAssign x t e -> AnnAssign <$> mapExprM f x <*> mapExprM f t <*> traverse (mapExprM f) e + For x iter body orelse -> For <$> mapExprM f x <*> mapExprM f iter <*> mapExprStatementsM f body <*> mapExprStatementsM f orelse + AsyncFor x iter body orelse -> AsyncFor <$> mapExprM f x <*> mapExprM f iter <*> mapExprStatementsM f body <*> mapExprStatementsM f orelse + While e body orelse -> While <$> mapExprM f e <*> mapExprStatementsM f body <*> mapExprStatementsM f orelse + If e body orelse -> If <$> mapExprM f e <*> mapExprStatementsM f body <*> mapExprStatementsM f orelse + With withitems body -> With <$> mapM (\(e1, e2) -> (,) <$> mapExprM f e1 <*> traverse (mapExprM f) e2) withitems <*> mapExprStatementsM f body + AsyncWith withitems body -> AsyncWith <$> mapM (\(e1, e2) -> (,) <$> mapExprM f e1 <*> traverse (mapExprM f) e2) withitems <*> mapExprStatementsM f body + Raise e1 e2 -> Raise <$> traverse (mapExprM f) e1 <*> traverse (mapExprM f) e2 + Try body handlers orelse finalbody -> Try <$> mapExprStatementsM f body <*> mapM (mapExprExceptHanderM f) handlers <*> mapExprStatementsM f orelse <*> mapExprStatementsM f finalbody + Assert e1 e2 -> Assert <$> mapExprM f e1 <*> traverse (mapExprM f) e2 + Import aliases -> return $ Import aliases + ImportFrom xs aliases -> return $ ImportFrom xs aliases + Global xs -> return $ Global xs + Nonlocal xs -> return $ Nonlocal xs + Expr' e -> Expr' <$> mapExprM f e + Pass -> return Pass + Break -> return Break + Continue -> return Continue + +mapExprStatementsM :: Monad m => (Expr' -> m Expr') -> [Statement'] -> m [Statement'] +mapExprStatementsM f = mapM (mapExprStatementM f) + +mapExprProgramM :: Monad m => (Expr' -> m Expr') -> Program -> m Program +mapExprProgramM = mapExprStatementsM + +mapExprProgram :: (Expr' -> Expr') -> Program -> Program +mapExprProgram f = runIdentity . mapExprStatementsM (return . f) diff --git a/src/Jikka/Python/Parse.hs b/src/Jikka/Python/Parse.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Python/Parse.hs @@ -0,0 +1,17 @@ +{-# LANGUAGE FlexibleContexts #-} + +module Jikka.Python.Parse + ( run, + ) +where + +import Data.Text (Text, unpack) +import Jikka.Common.Error +import Jikka.Python.Language.Expr (Program) +import qualified Jikka.Python.Parse.Alex as L +import qualified Jikka.Python.Parse.Happy as P + +run :: MonadError Error m => FilePath -> Text -> m Program +run _ input = do + tokens <- L.run $ unpack input + P.run tokens diff --git a/src/Jikka/Python/Parse/Alex.x b/src/Jikka/Python/Parse/Alex.x new file mode 100644 --- /dev/null +++ b/src/Jikka/Python/Parse/Alex.x @@ -0,0 +1,251 @@ +{ +-- vim: filetype=haskell +{-# LANGUAGE FlexibleContexts #-} + +-- | +-- Module : Jikka.Core.Parse.Alex +-- Description : tokenizes the code of the standard Python with Alex. +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- * TODO: tokenize float literals +-- * TODO: tokenize string literals +module Jikka.Python.Parse.Alex + ( run + ) where + +import Data.Char (chr, isHexDigit, isOctDigit) +import Jikka.Common.Error +import Jikka.Common.Location +import Jikka.Common.Parse.JoinLines (joinLinesWithParens, removeEmptyLines) +import Jikka.Common.Parse.OffsideRule (insertIndents) +import Jikka.Python.Parse.Token +} + +%wrapper "monad" + +$space = [\ ] +$tab = [\t] + +$alpha = [A-Z a-z] +$alnum = [0-9 A-Z a-z] +$doublequote = ["] +$backslash = [\\] +@nl = "\n" | "\r\n" + +$digit = [0-9] +$nonzerodigit = [1-9] +$bindigit = [0-1] +$octdigit = [0-7] +$hexdigit = [0-9a-fA-F] + +$shortstringchar_single = [^ \\ \r \n '] +$shortstringchar_double = [^ \\ \r \n '] +@stringescapeseq = $backslash . + +tokens :- + + $space + ; + "#" .* ; + $backslash @nl ; + @nl { tok Newline } + [\n\r] { tok Newline } + + "None" { tok None } + "True" { tok (Bool True) } + "False" { tok (Bool False) } + + "0" ("_" ? "0") * { tok' parseInt } + $nonzerodigit ("_" ? $digit) * { tok' parseInt } + "0" [bB] ("_" ? $bindigit) + { tok' parseInt } + "0" [oO] ("_" ? $octdigit) + { tok' parseInt } + "0" [xX] ("_" ? $hexdigit) + { tok' parseInt } + + "'" ($shortstringchar_single | @stringescapeseq) * "'" { tok'' parseString } + $doublequote ($shortstringchar_double | @stringescapeseq) * $doublequote { tok'' parseString } + + "def" { tok Def } + "if" { tok If } + "elif" { tok Elif } + "else" { tok Else } + "for" { tok For } + "in" { tok In } + "assert" { tok Assert } + "return" { tok Return } + "lambda" { tok Lambda } + + -- punctuations + "->" { tok Arrow } + ":" { tok Colon } + ";" { tok Semicolon } + "," { tok Comma } + "." { tok Dot } + "=" { tok Equal } + "_" { tok Underscore } + + -- parens + "{" { tok OpenBrace } + "[" { tok OpenBracket } + "(" { tok OpenParen } + "}" { tok CloseBrace } + "]" { tok CloseBracket } + ")" { tok CloseParen } + + -- special operators + "-" { tok MinusOp } + "*" { tok MulOp } + "**" { tok PowOp } + + -- expr operators + "+" { tok PlusOp } + "//" { tok (DivModOp FloorDiv) } + "/" { tok (DivModOp Div) } + "%" { tok (DivModOp FloorMod) } + "&" { tok BitAndOp } + "|" { tok BitOrOp } + "^" { tok BitXorOp } + "~" { tok BitNotOp } + "<<" { tok BitLShiftOp } + ">>" { tok BitRShiftOp } + ">" { tok (CmpOp GreaterThan) } + "<" { tok (CmpOp LessThan) } + "<=" { tok (CmpOp LessEqual) } + ">=" { tok (CmpOp GreaterEqual) } + "==" { tok (CmpOp DoubleEqual) } + "!=" { tok (CmpOp NotEqual) } + "and" { tok AndOp } + "or" { tok OrOp } + "not" { tok NotOp } + "@" { tok AtOp } + ":=" { tok WalrusOp } + + -- assignment operators + "+=" { tok (AugOp AugAdd) } + "-=" { tok (AugOp AugSub) } + "*=" { tok (AugOp AugMul) } + "@=" { tok (AugOp AugAt) } + "/=" { tok (AugOp AugDiv) } + "//=" { tok (AugOp AugFloorDiv) } + "%=" { tok (AugOp AugFloorMod) } + "%^=" { tok (AugOp AugCeilMod) } + "**=" { tok (AugOp AugPow) } + "<<=" { tok (AugOp AugBitRShift) } + ">>=" { tok (AugOp AugBitLShift) } + "&=" { tok (AugOp AugBitAnd) } + "^=" { tok (AugOp AugBitXor) } + "|=" { tok (AugOp AugBitOr) } + + -- additional operators + "/^" { tok (DivModOp CeilDiv) } + "%^" { tok (DivModOp CeilMod) } + "?" { tok MaxOp } + "implies" { tok ImpliesOp } + "/^=" { tok (AugOp AugCeilDiv) } + "?=" { tok (AugOp AugMax) } + + -- Python reserved + "as" { tok As } + "async" { tok Async } + "await" { tok Await } + "break" { tok Break } + "class" { tok Class } + "continue" { tok Continue } + "del" { tok Del } + "except" { tok Except } + "finally" { tok Finally } + "from" { tok From } + "global" { tok Global } + "import" { tok Import } + "is" { tok Is } + "nonlocal" { tok Nonlocal } + "pass" { tok Pass } + "raise" { tok Raise } + "try" { tok Try } + "while" { tok While } + "with" { tok With } + "yield" { tok Yield } + + -- identifier + ($alpha | _) ($alnum | _) * { tok' Ident } + + -- catch error + . { skip' } +{ +type Token'' = Either Error Token' + +alexEOF :: Alex (Maybe Token'') +alexEOF = return Nothing + +tok'' :: (Loc -> String -> Token'') -> AlexAction (Maybe Token'') +tok'' f (AlexPn _ line column, _, _, s) n = return . Just $ f loc (take n s) where + loc = Loc + { line = line + , column = column + , width = n + } + +tok' :: (String -> Token) -> AlexAction (Maybe Token'') +tok' f = tok'' (\loc s -> Right (WithLoc loc (f s))) + +tok :: Token -> AlexAction (Maybe Token'') +tok token = tok' (const token) + +parseInt :: String -> Token +parseInt s' = Int $ case filter (/= '_') s' of + '0' : 'b' : s -> foldl (\acc c -> acc * 2 + read [c]) 0 (reverse s) + '0' : 'B' : s -> foldl (\acc c -> acc * 2 + read [c]) 0 (reverse s) + s@('0' : 'o' : _) -> read s + s@('0' : 'O' : _) -> read s + s@('0' : 'x' : _) -> read s + s@('0' : 'X' : _) -> read s + s -> read s + +parseString :: Loc -> String -> Token'' +parseString loc s = WithLoc loc . String <$> go (tail (init s)) where + go "" = Right "" + go ('\\' : s) = case s of + [] -> throwInternalErrorAt loc "invalid escape sequence" + 'a' : s -> ('\a' :) <$> go s + 'b' : s -> ('\b' :) <$> go s + 'f' : s -> ('\f' :) <$> go s + 'n' : s -> ('\n' :) <$> go s + 'r' : s -> ('\r' :) <$> go s + 't' : s -> ('\t' :) <$> go s + 'v' : s -> ('\v' :) <$> go s + o1 : o2 : o3 : s | isOctDigit o1 && isOctDigit o2 && isOctDigit o3 -> (chr (read ("0o" ++ [o1, o2, o3])) :) <$> go s + o1 : o2 : s | isOctDigit o1 && isOctDigit o2 -> (chr (read ("0o" ++ [o1, o2])) :) <$> go s + o1 : s | isOctDigit o1 -> (chr (read ("0o" ++ [o1])) :) <$> go s + 'x' : h1 : h2 : s | isHexDigit h1 && isHexDigit h2 -> (chr (read ("0x" ++ [h1, h2])) :) <$> go s + 'x' : _ -> throwLexicalErrorAt loc "truncated \\xXX escape" + c : s -> (c :) <$> go s + go (c : s) = (c :) <$> go s + +skip' :: AlexAction (Maybe Token'') +skip' (AlexPn _ line column, _, _, s) n = return (Just (Left err)) where + loc = Loc line column n + msg = show (take n s) ++ " is not a acceptable character" + err = lexicalErrorAt loc msg + +unfoldM :: Monad m => m (Maybe a) -> m [a] +unfoldM f = do + x <- f + case x of + Nothing -> return [] + Just x -> (x :) <$> unfoldM f + +run :: MonadError Error m => String -> m [Token'] +run input = wrapError' "Jikka.Python.Parse.Alex failed" $ do + tokens <- case runAlex input (unfoldM alexMonadScan) of + Left err -> throwInternalError $ "Alex says: " ++ err + Right tokens -> return tokens + tokens <- reportErrors tokens + tokens <- joinLinesWithParens (`elem` [OpenParen, OpenBracket, OpenBrace]) (`elem` [CloseParen, CloseBracket, CloseBrace]) (== Newline) tokens + tokens <- return $ removeEmptyLines (== Newline) tokens + tokens <- insertIndents Indent Dedent (== Newline) tokens + return tokens +} diff --git a/src/Jikka/Python/Parse/Happy.y b/src/Jikka/Python/Parse/Happy.y new file mode 100644 --- /dev/null +++ b/src/Jikka/Python/Parse/Happy.y @@ -0,0 +1,553 @@ +{ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- vim: filetype=haskell + +-- | +-- Module : Jikka.Core.Parse.Happy +-- Description : parses the code of the standard Python with Happy. +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Python.Parse.Happy (run) where + +import Control.Arrow (first) +import Data.Functor (($>)) +import Data.List (intercalate) +import Data.Maybe (fromMaybe, isJust) +import qualified Data.Map.Strict as M +import Jikka.Common.Error +import Jikka.Common.Location +import Jikka.Python.Language.Expr +import qualified Jikka.Python.Parse.Token as L +} + +%name runHappy +%tokentype { WithLoc L.Token } +%monad { Either Error } +%error { happyErrorExpList } +%errorhandlertype explist + +%token + -- literals + "None" { WithLoc _ L.None } + INTEGER { WithLoc _ (L.Int _) } + BOOLEAN { WithLoc _ (L.Bool _) } + STRING { WithLoc _ (L.String _) } + BYTES { WithLoc _ (L.Bytes _) } + FLOAT { WithLoc _ (L.Float _) } + IMAGINARY { WithLoc _ (L.Imaginary _) } + + -- keywords + "def" { WithLoc _ L.Def } + "if" { WithLoc _ L.If } + "elif" { WithLoc _ L.Elif } + "else" { WithLoc _ L.Else } + "for" { WithLoc _ L.For } + "in" { WithLoc _ L.In } + "assert" { WithLoc _ L.Assert } + "return" { WithLoc _ L.Return } + "lambda" { WithLoc _ L.Lambda } + + -- punctuations + "->" { WithLoc _ L.Arrow } + ":" { WithLoc _ L.Colon } + ";" { WithLoc _ L.Semicolon } + "," { WithLoc _ L.Comma } + "." { WithLoc _ L.Dot } + "=" { WithLoc _ L.Equal } + "_" { WithLoc _ L.Underscore } + + -- parens + "[" { WithLoc _ L.OpenBracket } + "(" { WithLoc _ L.OpenParen } + "{" { WithLoc _ L.OpenBrace } + "]" { WithLoc _ L.CloseBracket } + ")" { WithLoc _ L.CloseParen } + "}" { WithLoc _ L.CloseBrace } + + -- identifier + IDENT { WithLoc _ (L.Ident _) } + + -- operator + ":=" { WithLoc _ L.WalrusOp } + "implies" { WithLoc _ L.ImpliesOp } + "or" { WithLoc _ L.OrOp } + "and" { WithLoc _ L.AndOp } + "not" { WithLoc _ L.NotOp } + COMP_OPERATOR { WithLoc _ (L.CmpOp $$) } + "?" { WithLoc _ L.MaxOp } + "|" { WithLoc _ L.BitOrOp } + "^" { WithLoc _ L.BitXorOp } + "&" { WithLoc _ L.BitAndOp } + "<<" { WithLoc _ L.BitLShiftOp } + ">>" { WithLoc _ L.BitRShiftOp } + "+" { WithLoc _ L.PlusOp } + "-" { WithLoc _ L.MinusOp } + "*" { WithLoc _ L.MulOp } + DIVMOD_OPERATOR { WithLoc _ (L.DivModOp $$) } + "~" { WithLoc _ L.BitNotOp } + "**" { WithLoc _ L.PowOp } + "@" { WithLoc _ L.AtOp } + AUGOP { WithLoc _ (L.AugOp $$) } + + -- indent + NEWLINE { WithLoc _ L.Newline } + INDENT { WithLoc _ L.Indent } + DEDENT { WithLoc _ L.Dedent } + + -- reserved + "as" { WithLoc _ L.As } + "async" { WithLoc _ L.Async } + "await" { WithLoc _ L.Await } + "break" { WithLoc _ L.Break } + "class" { WithLoc _ L.Class } + "continue" { WithLoc _ L.Continue } + "del" { WithLoc _ L.Del } + "except" { WithLoc _ L.Except } + "finally" { WithLoc _ L.Finally } + "from" { WithLoc _ L.From } + "global" { WithLoc _ L.Global } + "import" { WithLoc _ L.Import } + "is" { WithLoc _ L.Is } + "nonlocal" { WithLoc _ L.Nonlocal } + "pass" { WithLoc _ L.Pass } + "raise" { WithLoc _ L.Raise } + "try" { WithLoc _ L.Try } + "while" { WithLoc _ L.While } + "with" { WithLoc _ L.With } + "yield" { WithLoc _ L.Yield } +%% + +file_input :: { [Statement'] } + : {- empty -} { [] } + | file_input NEWLINE { $1 } + | file_input statement { $1 ++ $2 } + +-- utilities +opt(p) -- :: { Maybe a } + : {- empty -} { Nothing } + | p { Just $1 } +rev_list1(p) -- :: { [a] } + : p { [$1] } + | rev_list1(p) p { $2 : $1 } +list1(p) -- :: { [a] } + : rev_list1(p) { reverse $1 } +list(p) -- :: { [a] } + : {- empty -} { [] } + | list1(p) { $1 } +rev_sep1(p, q) -- :: { [a] } + : p { [$1] } + | rev_sep1(p, q) q p { $3 : $1 } +sep1(p, q) -- :: { [a] } + : rev_sep1(p, q) { reverse $1 } +sep1opt(p, q) -- :: { [a] } + : rev_sep1(p, q) opt(q) { reverse $1 } +fst(p, q) + : p q { $1 } +snd(p, q) + : p q { $2 } +both(p, q) + : p q { ($1, $2) } + +-- 6.2 Atoms +atom :: { Expr' } + : identifier { $1 @> Name $1 } + | literal { Constant `fmap` $1 } + | enclosure { $1 } +enclosure :: { Expr' } + : parenth_form { $1 } + | list_display { $1 } + +-- 6.2.1 Identifiers +identifier :: { Ident' } + : IDENT { let (L.Ident x) = value $1 in $1 @> Ident x } + | "_" { $1 @> Ident "_" } + +-- 6.2.2 Literals +literal :: { WithLoc Constant } + : "None" { $1 @> ConstNone } + | INTEGER { let (L.Int n) = value $1 in $1 @> ConstInt n } + | BOOLEAN { let (L.Bool p) = value $1 in $1 @> ConstBool p } + | STRING { let (L.String s) = value $1 in $1 @> ConstString s } + | BYTES { let (L.Bytes s) = value $1 in $1 @> ConstBytes s } + | FLOAT { let (L.Float x) = value $1 in $1 @> ConstFloat x } + | IMAGINARY { let (L.Imaginary y) = value $1 in $1 @> ConstImaginary y } + +-- 6.2.3 Parenthesized forms +parenth_form :: { Expr' } + : "(" ")" { $1 @> Tuple [] } + | "(" expression_list ")" { uncurry fromExprList $2 } + +-- 6.2.4 Displays for lists, sets and dictionaries +comprehension :: { (Expr', [Comprehension]) } + : expression comp_for { ($1, $2) } +comp_for :: { [Comprehension] } + : "for" identifier "in" implies_test opt(comp_if) { [Comprehension ($2 @> Name $2) $4 $5] } +comp_if :: { Expr' } + : "if" expression_nocond { $2 } + +-- 6.2.5 List displays +list_display :: { Expr' } + : "[" "]" { $1 @> List [] } + | "[" expression_list "]" { $1 @> List (fst $2) } + | "[" comprehension "]" { $1 @> uncurry ListComp $2 } + +-- 6.2.9. Yield expressions +yield_expression :: { Expr' } + : "yield" opt(expression_list) { $1 @> Yield (uncurry fromExprList `fmap` $2) } + | "yield" "from" expression { $1 @> YieldFrom $3 } + +-- 6.3 Primaries +primary :: { Expr' } + : atom { $1 } + | attributeref { $1 } + | subscription { $1 } + | slicing { $1 } + | call { $1 } + +-- 6.3.1. Attribute references +attributeref :: { Expr' } + : primary "." identifier { $1 @> Attribute $1 $3 } + +-- 6.3.2. Subscriptions +subscription :: { Expr' } + : primary "[" expression_list "]" { $1 @> Subscript $1 (uncurry fromExprList $3) } + +-- 6.3.3. Slicings +slicing :: { Expr' } + : primary "[" opt(expression) ":" opt(expression) opt(snd(":", expression)) "]" { $1 @> Subscript $1 ($2 @> Slice $3 $5 $6) } + +-- 6.3.4. Calls +call :: { Expr' } + : primary "(" ")" { $1 @> Call $1 [] [] } + | primary "(" starred_list ")" { $1 @> Call $1 (fst $3) [] } + | primary "(" comprehension ")" { $1 @> Call $1 [$2 @> uncurry GeneratorExp $3] [] } + +-- 6.5. The power operator +power :: { Expr' } + : primary { $1 } + | primary "**" u_expr { $1 @> BinOp $1 Pow $3 } + +-- 6.6. Unary arithmetic and bitwise operations +u_expr :: { Expr' } + : power { $1 } + | "-" u_expr { $1 @> UnaryOp USub $2 } + | "+" u_expr { $1 @> UnaryOp UAdd $2 } + | "~" u_expr { $1 @> UnaryOp Invert $2 } + +-- 6.7. Binary arithmetic operations +m_expr :: { Expr' } + : u_expr { $1 } + | m_expr "*" u_expr { $1 @> BinOp $1 Mult $3 } + | m_expr "@" u_expr { $1 @> BinOp $1 MatMult $3 } + | m_expr DIVMOD_OPERATOR u_expr { $1 @> BinOp $1 (fromDivModOp $2) $3 } +a_expr :: { Expr' } + : m_expr { $1 } + | a_expr "+" m_expr { $1 @> BinOp $1 Add $3 } + | a_expr "-" m_expr { $1 @> BinOp $1 Sub $3 } + +-- 6.8. Shifting operations +shift_expr :: { Expr' } + : a_expr { $1 } + | shift_expr "<<" a_expr { $1 @> BinOp $1 BitLShift $3 } + | shift_expr ">>" a_expr { $1 @> BinOp $1 BitRShift $3 } + +-- 6.9. Binary bitwise operations +and_expr :: { Expr' } + : shift_expr { $1 } + | and_expr "&" shift_expr { $1 @> BinOp $1 BitAnd $3 } +xor_expr :: { Expr' } + : and_expr { $1 } + | xor_expr "^" and_expr { $1 @> BinOp $1 BitXor $3 } +or_expr :: { Expr' } + : xor_expr { $1 } + | or_expr "|" xor_expr { $1 @> BinOp $1 BitOr $3 } + +-- Extra.1. Min and max operations +min_expr :: { Expr' } + : or_expr { $1 } + | min_expr " BinOp $1 Min $3 } + | min_expr ">?" or_expr { $1 @> BinOp $1 Max $3 } + +-- 6.10. Comparisons +comparison :: { (Expr', [(CmpOp, Expr')]) } + : min_expr { ($1, []) } + | comparison comp_operator min_expr { let (e1, e2) = $1 in (e1, e2 ++ [($2, $3)]) } +comp_operator :: { CmpOp } + : COMP_OPERATOR { fromCmpOp $1 } + | "is" { Is } + | "is" "not" { IsNot } + | "in" { In } + | "not" "in" { NotIn } + +-- 6.11. Boolean operations +not_test :: { Expr' } + : comparison { convertCompare $1 } + | "not" not_test { $1 @> UnaryOp Not $2 } +and_test :: { Expr' } + : not_test { $1 } + | and_test "and" not_test { $1 @> BoolOp $1 And $3 } +or_test :: { Expr' } + : and_test { $1 } + | or_test "or" and_test { $1 @> BoolOp $1 Or $3 } + +-- Extra.2. Implication operation +implies_test :: { Expr' } + : or_test { $1 } + | or_test "implies" implies_test { $1 @> BoolOp $1 Implies $3 } + +-- 6.13. Conditional expressions +conditional_expression :: { Expr' } + : implies_test { $1 } + | implies_test "if" implies_test "else" expression { $1 @> IfExp $3 $1 $5 } +expression :: { Expr' } + : conditional_expression { $1 } + | lambda_expr { $1 } +expression_nocond :: { Expr' } + : implies_test { $1 } + | lambda_expr_nocond { $1 } + +-- 6.14. Lambda +lambda_expr :: { Expr' } + : "lambda" ":" expression { $1 @> Lambda emptyArguments $3} + | "lambda" sep1opt(identifier, ",") ":" expression { $1 @> Lambda (convertArguments' (reverse $2)) $4} +lambda_expr_nocond :: { Expr' } + : "lambda" ":" expression_nocond { $1 @> Lambda emptyArguments $3} + | "lambda" sep1opt(identifier, ",") ":" expression_nocond { $1 @> Lambda (convertArguments' (reverse $2)) $4} + +-- 6.15. Expression lists +expression_list :: { ([Expr'], Bool) } + : expression opt(",") { ([$1], isJust $2) } + | expression "," expression_list { first ($1 :) $3 } +starred_list :: { ([Expr'], Bool) } + : starred_item opt(",") { ([$1], isJust $2) } + | starred_item "," starred_list { first ($1 :) $3 } +starred_item :: { Expr' } + : expression { $1 } + | "*" min_expr { $1 @> Starred $2 } + +-- 7. Simple statements +simple_stmt :: { Statement' } + : expression_stmt { $1 } + | assert_stmt { $1 } + | assignment_stmt { $1 } + | augmented_assignment_stmt { $1 } + | annotated_assignment_stmt { $1 } + | pass_stmt { $1 } + | del_stmt { $1 } + | return_stmt { $1 } + | yield_stmt { $1 } + | raise_stmt { $1 } + | break_stmt { $1 } + | continue_stmt { $1 } + | import_stmt { $1 } + | global_stmt { $1 } + | nonlocal_stmt { $1 } + +-- 7.1. Expression statements +expression_stmt :: { Statement' } + : expression { $1 @> Expr' $1 } + +-- 7.2. Assignment statements +assignment_stmt :: { Statement' } + : expression_list "=" expression { convertAssign $1 $3 } + +-- 7.2.1. Augmented assignment statements +augmented_assignment_stmt :: { Statement' } + : augtarget AUGOP expression_list { $1 @> AugAssign $1 (fromAugOp $2) (uncurry fromExprList $3) } +augtarget :: { Target' } + : identifier { $1 @> Name $1 } + | attributeref { $1 } + | subscription { $1 } + | slicing { $1 } + +-- 7.2.2. Annotated assignment statements +annotated_assignment_stmt :: { Statement' } + : augtarget ":" expression opt(snd("=", expression)) { $1 @> AnnAssign $1 $3 $4 } + +-- 7.3. The assert statement +assert_stmt :: { Statement' } + : "assert" expression opt(snd(",", expression)) { $1 @> Assert $2 $3 } + +-- 7.4. The pass statement +pass_stmt :: { Statement' } + : "pass" { $1 @> Pass } + +-- 7.5. The del statement¶ +del_stmt :: { Statement' } + : "del" expression_list { $1 @> Delete (fst $2) } + +-- 7.6. The return statement +return_stmt :: { Statement' } + : "return" opt(expression_list) { $1 @> Return (uncurry fromExprList `fmap` $2) } + +-- 7.7. The yield statement +yield_stmt :: { Statement' } + : yield_expression { $1 @> Expr' $1 } + +-- 7.8. The raise statement +raise_stmt :: { Statement' } + : "raise" opt(expression) { $1 @> Raise $2 Nothing } + | "raise" expression "from" expression { $1 @> Raise (Just $2) (Just $4) } + +-- 7.9. The break statement +break_stmt :: { Statement' } + : "break" { $1 @> Break } + +-- 7.10. The continue statement +continue_stmt :: { Statement' } + : "continue" { $1 @> Continue } + +-- 7.11. The import statement +import_stmt :: { Statement' } + : "import" module_ opt(snd("as", identifier)) list(both(snd(",", module_), opt(snd("as", identifier)))) { $1 @> Import [] } + | "from" module_ "import" "*" { $1 @> ImportFrom [] [] } +module_ :: { [Ident'] } + : sep1(identifier, ".") { $1 } + +-- 7.12. The global statement +global_stmt :: { Statement' } + : "global" sep1(identifier, ",") { $1 @> Global $2 } + +-- 7.13. The nonlocal statement +nonlocal_stmt :: { Statement' } + : "nonlocal" sep1(identifier, ",") { $1 @> Nonlocal $2 } + +-- 8. Compound statements +compound_stmt :: { Statement' } + : if_stmt { $1 } + | while_stmt { $1 } + | for_stmt { $1 } + | funcdef { $1 } +suite :: { [Statement'] } + : stmt_list NEWLINE { $1 } + | NEWLINE INDENT list1(statement) DEDENT { concat $3 } +statement :: { [Statement'] } + : stmt_list NEWLINE { $1 } + | compound_stmt { [$1] } +stmt_list :: { [Statement'] } + : simple_stmt opt(";") { [$1] } + | simple_stmt ";" stmt_list { $1 : $3 } + +-- 8.1. The if statement +if_stmt :: { Statement' } + : "if" expression ":" suite list(both(both("elif", expression), snd(":", suite))) opt(snd("else", snd(":", suite))) { convertIfElse $1 $2 $4 $5 $6 } + +-- 8.2. The while statement +while_stmt :: { Statement' } + : "while" expression ":" suite opt(snd("else", snd(":", suite))) { $1 @> While $2 $4 (fromMaybe [] $5) } + +-- 8.3. The for statement +for_stmt :: { Statement' } + : "for" identifier "in" expression_list ":" suite opt(snd("else", snd(":", suite))) { $1 @> For ($2 @> Name $2) (uncurry fromExprList $4) $6 (fromMaybe [] $7) } + +-- 8.6. Function definitions +funcdef :: { Statement' } + : list(decorator) "def" funcname "(" opt(parameter_list) ")" opt(snd("->", expression)) ":" suite { $2 @> FunctionDef $3 (convertArguments $5) $9 $1 $7 } +decorator :: { Decorator' } + : "@" expression NEWLINE { $2 } +parameter_list :: { [Arg] } + : parameter opt(",") { [$1] } + | parameter "," parameter_list { $1 : $3 } +parameter :: { Arg } + : identifier opt(snd(":", expression)) { ($1, $2) } +funcname :: { Ident' } + : identifier { $1 } + +{ +(@>) :: WithLoc a -> b -> WithLoc b +(@>) = ($>) + +fromExprList :: [Expr'] -> Bool -> Expr' +fromExprList [] _ = bug "empty list for fromExprList" +fromExprList [e] False = e +fromExprList es@(e : _) _ = e $> Tuple es + +fromCmpOp :: L.CmpOp -> CmpOp +fromCmpOp = \case + L.DoubleEqual -> Eq' + L.NotEqual -> NotEq + L.LessThan -> Lt + L.LessEqual -> LtE + L.GreaterThan -> Gt + L.GreaterEqual -> GtE + +fromDivModOp :: L.DivModOp -> Operator +fromDivModOp = \case + L.Div -> Div + L.FloorDiv -> FloorDiv + L.FloorMod -> FloorMod + L.CeilDiv -> CeilDiv + L.CeilMod -> CeilMod + +fromAugOp :: L.AugOp -> Operator +fromAugOp = \case + L.AugAdd -> Add + L.AugSub -> Sub + L.AugMul -> Mult + L.AugAt -> MatMult + L.AugDiv -> Div + L.AugFloorDiv -> FloorDiv + L.AugFloorMod -> FloorMod + L.AugCeilDiv -> CeilDiv + L.AugCeilMod -> CeilMod + L.AugPow -> Pow + L.AugBitRShift -> BitRShift + L.AugBitLShift -> BitLShift + L.AugBitAnd -> BitAnd + L.AugBitXor -> BitXor + L.AugBitOr -> BitOr + L.AugMin -> Min + L.AugMax -> Max + +convertArguments :: Maybe [Arg] -> Arguments +convertArguments args = emptyArguments { argsArgs = fromMaybe [] args } + +convertArguments' :: [Ident'] -> Arguments +convertArguments' args = emptyArguments { argsArgs = map (\x -> (x, Nothing)) args } + +convertIfElse :: WithLoc a -> Expr' -> [Statement'] -> [((WithLoc a, Expr'), [Statement'])] -> Maybe [Statement'] -> Statement' +convertIfElse head cond body elifs orelse = head @> If cond body cont where + cont = case elifs of + [] -> fromMaybe [] orelse + ((head', cond'), body') : elifs -> [convertIfElse head' cond' body' elifs orelse] + +convertCompare :: (Expr', [(CmpOp, Expr')]) -> Expr' +convertCompare (e, []) = e +convertCompare (e, ops) = e $> Compare e ops + +convertAssign :: ([Expr'], Bool) -> Expr' -> Statement' +convertAssign ([], _) _ = bug "empty targets for convertAssign" +convertAssign (xs, comma) e = head xs $> Assign [fromExprList xs comma] e + +happyErrorExpList :: ([WithLoc L.Token], [String]) -> Either Error a +happyErrorExpList (tokens, expected) = Left err where + err :: Error + err = WithGroup SyntaxError (withLocation tokens (Error msg)) + withLocation :: [WithLoc L.Token] -> Error -> Error + withLocation [] err = err + withLocation (token : _) err = WithLocation(loc token) err + msg :: String + msg = tok tokens ++ " is got, but " ++ exp expected ++ " expected" + tok :: [WithLoc L.Token] -> String + tok [] = "EOF" + tok (token : _) = wrap . show $ value token + exp :: [String] -> String + exp [] = "EOF is" + exp [item] = wrap item ++ " is" + exp items = intercalate ", " (map wrap $ init items) ++ ", or " ++ (wrap $ last items) ++ " are" + wrap :: String -> String + wrap ('\'' : s) = '`' : s + wrap s = "`" ++ s ++ "'" + +run :: MonadError Error m => [WithLoc L.Token] -> m Program +run tokens = wrapError' "Jikka.Python.Parse.Happy.run failed" $ do + case runHappy tokens of + Left err -> throwError err + Right stmts -> return stmts +} diff --git a/src/Jikka/Python/Parse/Token.hs b/src/Jikka/Python/Parse/Token.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/Python/Parse/Token.hs @@ -0,0 +1,137 @@ +-- | +-- Module : Jikka.Core.Parse.Token +-- Description : defines tokens of the standard Python. / 標準の Python の字句要素を定義します。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.Python.Parse.Token where + +import Data.Int (Int8) +import Jikka.Common.Location + +data CmpOp + = DoubleEqual + | NotEqual + | LessThan + | LessEqual + | GreaterThan + | GreaterEqual + deriving (Eq, Ord, Show, Read) + +data DivModOp + = Div + | FloorDiv + | FloorMod + | CeilDiv + | CeilMod + deriving (Eq, Ord, Show, Read) + +data AugOp + = AugAdd + | AugSub + | AugMul + | AugAt + | AugDiv + | AugFloorDiv + | AugFloorMod + | AugCeilDiv + | AugCeilMod + | AugPow + | AugBitRShift + | AugBitLShift + | AugBitAnd + | AugBitXor + | AugBitOr + | AugMin + | AugMax + deriving (Eq, Ord, Show, Read) + +-- | We don't have to classify tokens in detail, but it's convenient for testing and debugging. +data Token + = -- literals + None + | Int Integer + | Bool Bool + | String String + | Bytes [Int8] + | Float Double + | Imaginary Double + | -- keywords + Def + | If + | Elif + | Else + | For + | In + | Assert + | Return + | Lambda + | -- punctuations + Arrow + | Colon + | Semicolon + | Comma + | Dot + | Equal + | Underscore + | -- parens + OpenBrace + | OpenBracket + | OpenParen + | CloseBrace + | CloseBracket + | CloseParen + | -- identifier + Ident String + | -- operators + WalrusOp + | ImpliesOp + | OrOp + | AndOp + | NotOp + | CmpOp CmpOp + | MinOp + | MaxOp + | BitOrOp + | BitXorOp + | BitAndOp + | BitLShiftOp + | BitRShiftOp + | PlusOp + | MinusOp + | MulOp + | DivModOp DivModOp + | AtOp + | BitNotOp + | PowOp + | AugOp AugOp + | -- indent + Newline + | Indent + | Dedent + | -- reserved keywords + As + | Async + | Await + | Break + | Class + | Continue + | Del + | Except + | Finally + | From + | Global + | Import + | Is + | Nonlocal + | Pass + | Raise + | Try + | While + | With + | Yield + deriving (Eq, Ord, Show, Read) + +type Token' = WithLoc Token diff --git a/src/Jikka/RestrictedPython/Convert.hs b/src/Jikka/RestrictedPython/Convert.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Convert.hs @@ -0,0 +1,42 @@ +{-# LANGUAGE FlexibleContexts #-} + +module Jikka.RestrictedPython.Convert + ( run, + run', + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Common.IOFormat +import qualified Jikka.Core.Language.Expr as Y +import qualified Jikka.RestrictedPython.Convert.Alpha as Alpha +import qualified Jikka.RestrictedPython.Convert.DefaultMain as DefaultMain +import qualified Jikka.RestrictedPython.Convert.ParseMain as ParseMain +import qualified Jikka.RestrictedPython.Convert.RemoveUnbalancedIf as RemoveUnbalancedIf +import qualified Jikka.RestrictedPython.Convert.RemoveUnreachable as RemoveUnreachable +import qualified Jikka.RestrictedPython.Convert.ResolveBuiltin as ResolveBuiltin +import qualified Jikka.RestrictedPython.Convert.SplitLoops as SplitLoops +import qualified Jikka.RestrictedPython.Convert.ToCore as ToCore +import qualified Jikka.RestrictedPython.Convert.TypeInfer as TypeInfer +import qualified Jikka.RestrictedPython.Convert.UseAppend as UseAppend +import qualified Jikka.RestrictedPython.Language.Expr as X + +run' :: (MonadAlpha m, MonadError Error m) => X.Program -> m (X.Program, IOFormat) +run' prog = do + prog <- return $ RemoveUnreachable.run prog + prog <- return $ RemoveUnbalancedIf.run prog + prog <- UseAppend.run prog + prog <- ResolveBuiltin.run prog + prog <- Alpha.run prog + (format, prog) <- ParseMain.run prog -- Run ParseMain before type inference because main function has different semantics. + prog <- TypeInfer.run prog + prog <- SplitLoops.run prog + format <- maybe (DefaultMain.run prog) return format + return (prog, format) + +run :: (MonadAlpha m, MonadError Error m) => X.Program -> m (Y.Program, IOFormat) +run prog = do + (prog, format) <- run' prog + prog <- ToCore.run prog + return (prog, format) diff --git a/src/Jikka/RestrictedPython/Convert/Alpha.hs b/src/Jikka/RestrictedPython/Convert/Alpha.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Convert/Alpha.hs @@ -0,0 +1,308 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.RestrictedPython.Convert.Alpha +-- Description : does alpha conversion. / alpha 変換を行います。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.RestrictedPython.Convert.Alpha + ( run, + ) +where + +import Control.Monad.State.Strict +import Data.List (delete, intersect) +import Data.Maybe (isNothing) +import qualified Data.Set as S +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.RestrictedPython.Language.Builtin +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Lint +import Jikka.RestrictedPython.Language.Util +import Jikka.RestrictedPython.Language.VariableAnalysis + +data Env = Env + { currentMapping :: [(VarName, VarName)], + parentMappings :: [[(VarName, VarName)]] + } + deriving (Eq, Ord, Read, Show) + +initialEnv :: Env +initialEnv = + Env + { currentMapping = [], + parentMappings = [map (\x -> (x, x)) (S.toList builtinNames)] + } + +withToplevelScope :: (MonadError Error m, MonadState Env m) => m a -> m a +withToplevelScope f = do + env <- get + x <- catchError' $ withScope f + put env + liftEither x + +withScope :: (MonadError Error m, MonadState Env m) => m a -> m a +withScope f = do + modify' $ \env -> + env + { currentMapping = [], + parentMappings = currentMapping env : parentMappings env + } + x <- catchError' f + modify' $ \env -> + env + { currentMapping = head (parentMappings env), + parentMappings = tail (parentMappings env) + } + liftEither x + +-- | `renameLocalNew` renames given variables and record them to the `Env`. +renameLocalNew :: (MonadAlpha m, MonadState Env m) => VarName' -> m VarName' +renameLocalNew x = do + env <- get + case lookupLocalName x (env {currentMapping = []}) of + Just y -> return y + Nothing -> do + y <- genVarName x + when (unVarName (value' x) /= "_") $ do + put $ + env + { currentMapping = (value' x, value' y) : currentMapping env + } + return y + +-- | `renameShadow` renames given variables ignoring the current `Env` and record them to the `Env`. +renameShadow :: (MonadAlpha m, MonadState Env m) => VarName' -> m VarName' +renameShadow x = do + env <- get + y <- genVarName x + put $ + env + { currentMapping = (value' x, value' y) : currentMapping env + } + return y + +-- | `renameLocalCompletelyNew` throws errors when given variables already exists in environments. +renameLocalCompletelyNew :: (MonadAlpha m, MonadState Env m, MonadError Error m) => VarName' -> m VarName' +renameLocalCompletelyNew x = do + env <- get + case lookupLocalName x env of + Just _ -> throwSemanticErrorAt' (loc' x) $ "cannot redefine variable: " ++ unVarName (value' x) + Nothing -> renameLocalNew x + +-- | `renameToplevel` records given variables to the `Env` without actual renaming. +renameToplevel :: (MonadAlpha m, MonadState Env m, MonadError Error m) => VarName' -> m VarName' +renameToplevel x = do + env <- get + case lookupName x env of + Just _ -> do + let msg = + if value' x `S.member` builtinNames + then "cannot assign to builtin function: " ++ unVarName (value' x) + else "cannot redefine variable in toplevel: " ++ unVarName (value' x) + throwSemanticErrorAt' (loc' x) msg + Nothing -> do + when (unVarName (value' x) /= "_") $ do + put $ + env + { currentMapping = (value' x, value' x) : currentMapping env + } + return x + +-- | `renameToplevelArgument` always introduces a new variable. +renameToplevelArgument :: (MonadAlpha m, MonadState Env m, MonadError Error m) => VarName' -> m VarName' +renameToplevelArgument = renameShadow + +popRename :: (MonadState Env m, MonadError Error m) => VarName' -> m () +popRename x = + when (unVarName (value' x) /= "_") $ do + y <- lookupName' x + modify' $ \env -> env {currentMapping = delete (value' x, value' y) (currentMapping env)} + +lookupName :: VarName' -> Env -> Maybe VarName' +lookupName x env = lookupNameFromMappings x (currentMapping env : parentMappings env) + +lookupLocalName :: VarName' -> Env -> Maybe VarName' +lookupLocalName x env = lookupNameFromMappings x (reverse (drop 2 (reverse (currentMapping env : parentMappings env)))) + +lookupNameFromMappings :: VarName' -> [[(VarName, VarName)]] -> Maybe VarName' +lookupNameFromMappings _ [] = Nothing +lookupNameFromMappings x (mapping : mappings) = + case lookup (value' x) mapping of + Just y -> return $ WithLoc' (loc' x) y + Nothing -> lookupNameFromMappings x mappings + +lookupName' :: (MonadState Env m, MonadError Error m) => VarName' -> m VarName' +lookupName' x = do + env <- get + case lookupName x env of + Just y -> return y + Nothing -> throwSymbolErrorAt' (loc' x) $ "undefined identifier: " ++ unVarName (value' x) + +-- | `runAnnTarget` renames targets of annotated assignments. +runAnnTarget :: (MonadState Env m, MonadAlpha m, MonadError Error m) => Target' -> m Target' +runAnnTarget = runTargetGeneric renameLocalNew + +-- | `runForTarget` renames targets of for-loops. +runForTarget :: (MonadState Env m, MonadAlpha m, MonadError Error m) => Target' -> m Target' +runForTarget = runTargetGeneric renameLocalCompletelyNew + +-- | `runAugTarget` renames targets of augumented assignments. +runAugTarget :: (MonadState Env m, MonadAlpha m, MonadError Error m) => Target' -> m Target' +runAugTarget = runTargetGeneric lookupName' + +runTargetGeneric :: (MonadState Env m, MonadAlpha m, MonadError Error m) => (VarName' -> m VarName') -> Target' -> m Target' +runTargetGeneric f x = + WithLoc' (loc' x) <$> case value' x of + SubscriptTrg f index -> SubscriptTrg <$> runAugTarget f <*> runExpr index + NameTrg x -> NameTrg <$> f x + TupleTrg xs -> TupleTrg <$> mapM (runTargetGeneric f) xs + +popTarget :: (MonadState Env m, MonadError Error m) => Target' -> m () +popTarget (WithLoc' _ x) = case x of + SubscriptTrg _ _ -> return () + NameTrg x -> popRename x + TupleTrg xs -> mapM_ popTarget xs + +runExpr :: (MonadState Env m, MonadAlpha m, MonadError Error m) => Expr' -> m Expr' +runExpr e0 = + wrapAt' (loc' e0) $ + WithLoc' (loc' e0) <$> case value' e0 of + BoolOp e1 op e2 -> BoolOp <$> runExpr e1 <*> return op <*> runExpr e2 + BinOp e1 op e2 -> BinOp <$> runExpr e1 <*> return op <*> runExpr e2 + UnaryOp op e -> UnaryOp op <$> runExpr e + Lambda args body -> + withToplevelScope $ do + args <- forM args $ \(x, t) -> do + y <- renameLocalNew x + return (y, t) + body <- runExpr body + return $ Lambda args body + IfExp e1 e2 e3 -> IfExp <$> runExpr e1 <*> runExpr e2 <*> runExpr e3 + ListComp e (Comprehension x iter ifs) -> do + iter <- runExpr iter + y <- runAnnTarget x + ifs <- mapM runExpr ifs + e <- runExpr e + popTarget x + return $ ListComp e (Comprehension y iter ifs) + Compare e1 op e2 -> Compare <$> runExpr e1 <*> return op <*> runExpr e2 + Call f args -> Call <$> runExpr f <*> mapM runExpr args + Constant const -> return $ Constant const + Attribute e x -> Attribute <$> runExpr e <*> pure x + Subscript e1 e2 -> Subscript <$> runExpr e1 <*> runExpr e2 + Starred e -> Starred <$> runExpr e + Name x -> Name <$> lookupName' x + List t es -> List t <$> mapM runExpr es + Tuple es -> Tuple <$> mapM runExpr es + SubscriptSlice e from to step -> SubscriptSlice <$> runExpr e <*> mapM runExpr from <*> mapM runExpr to <*> mapM runExpr step + +runStatement :: (MonadState Env m, MonadAlpha m, MonadError Error m) => Statement -> m Statement +runStatement = \case + Return e -> Return <$> runExpr e + AugAssign x op e -> do + e <- runExpr e + x <- runAugTarget x + return $ AugAssign x op e + AnnAssign x t e -> do + e <- runExpr e -- visit e before x + x <- runAnnTarget x + return $ AnnAssign x t e + For x e body -> do + e <- runExpr e + withScope $ do + y <- runForTarget x + body <- runStatements body + return $ For y e body + If e body1 body2 -> do + e <- runExpr e + let (_, WriteList w1) = analyzeStatementsMin body1 + let (_, WriteList w2) = analyzeStatementsMin body2 + forM_ (w1 `intersect` w2) $ \x -> do + isLocallyUndefined <- isNothing . lookupLocalName (withoutLoc x) <$> get + when isLocallyUndefined $ do + renameLocalNew (withoutLoc x) -- introduce variables to the parent scope + return () + body1 <- withScope $ do + runStatements body1 + body2 <- withScope $ do + runStatements body2 + return $ If e body1 body2 + Assert e -> Assert <$> runExpr e + Expr' e -> Expr' <$> runExpr e + +runStatements :: (MonadState Env m, MonadAlpha m, MonadError Error m) => [Statement] -> m [Statement] +runStatements stmts = reportErrors =<< mapM (catchError' . runStatement) stmts + +runToplevelStatement :: (MonadState Env m, MonadAlpha m, MonadError Error m) => ToplevelStatement -> m ToplevelStatement +runToplevelStatement = \case + ToplevelAnnAssign x t e -> do + e <- runExpr e -- visit e before x + y <- renameToplevel x + return $ ToplevelAnnAssign y t e + ToplevelFunctionDef f args ret body -> do + g <- renameToplevel f + withToplevelScope $ do + args <- forM args $ \(x, t) -> do + y <- renameToplevelArgument x + return (y, t) + body <- runStatements body + return $ ToplevelFunctionDef g args ret body + ToplevelAssert e -> ToplevelAssert <$> runExpr e + +runProgram :: (MonadState Env m, MonadAlpha m, MonadError Error m) => Program -> m Program +runProgram prog = reportErrors =<< mapM (catchError' . runToplevelStatement) prog + +-- | `run` renames variables. +-- This assumes `doesntHaveAssignmentToBuiltin`. +-- +-- * This introduce a new name for each assignment if possible. +-- For example, the following +-- +-- > x = 21 +-- > x += x +-- > x = 42 +-- > x += x +-- > for _ in range(100): +-- > x = x + 1 +-- > x = x + 1 +-- +-- turns the following +-- +-- > x0 = 21 +-- > x1 += x0 +-- > x2 = 42 +-- > x3 += x2 +-- > for a4 in range(100): +-- > x3 = x3 + 1 +-- > x5 = x3 + 1 +-- +-- * This blames leaks of loop counters of for-statements, i.e. `doesntHaveLeakOfLoopCounters`. +-- For example, the followings is not allowed. +-- +-- > for i in range(10): +-- > a = 0 +-- > return a # error +-- +-- * This blames leaks of names from for-statements and if-statements at all. +-- For example, the followings are not allowed. +-- +-- > if True: +-- > a = 0 +-- > else: +-- > b = 1 +-- > return a # error +-- +-- > for i in range(10): +-- > a = 0 +-- > return a # error +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.RestrictedPython.Convert.Alpha" $ do + ensureDoesntHaveLeakOfLoopCounters prog + ensureDoesntHaveAssignmentToBuiltin prog + evalStateT (runProgram prog) initialEnv diff --git a/src/Jikka/RestrictedPython/Convert/DefaultMain.hs b/src/Jikka/RestrictedPython/Convert/DefaultMain.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Convert/DefaultMain.hs @@ -0,0 +1,87 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | +-- Module : Jikka.RestrictedPython.Convert.DefaultMain +-- Description : makes a default IO format based on types. / 型に基づくデフォルトの入出力フォーマットを作成します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.RestrictedPython.Convert.DefaultMain + ( run, + ) +where + +import Control.Arrow +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Common.IOFormat +import Jikka.RestrictedPython.Format (formatType) +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Util + +lookupSolve :: MonadError Error m => Program -> m (Maybe Loc, [(VarName', Type)], Type, [Statement]) +lookupSolve = \case + [] -> throwSymbolError "solve function is not defined" + ToplevelAnnAssign _ _ _ : stmts -> lookupSolve stmts + ToplevelFunctionDef f args ret body : stmts -> case value' f of + VarName "solve" -> return (loc' f, args, ret, body) + _ -> lookupSolve stmts + ToplevelAssert _ : stmts -> lookupSolve stmts + +makeInputFormatFromType :: (MonadAlpha m, MonadError Error m) => Type -> m (FormatTree, String) +makeInputFormatFromType = \case + IntTy -> do + x <- unVarName . value' <$> genVarName' + return (Exp (Var x), x) + ListTy t -> do + n <- unVarName . value' <$> genVarName' + i <- unVarName . value' <$> genVarName' + (body, x) <- makeInputFormatFromType t + body <- (`mapFormatTreeM` body) $ \case + Exp e -> return $ Exp (At e i) + format -> return format + return (Seq [Exp (Var n), Loop i (Var n) body], x) + t -> throwSemanticError $ "cannot read input of type: " ++ formatType t + +makeOutputFormatFromType' :: (MonadAlpha m, MonadError Error m) => Type -> m (FormatTree, String) +makeOutputFormatFromType' = \case + IntTy -> do + x <- unVarName . value' <$> genVarName' + return (Exp (Var x), x) + ListTy t -> do + i <- unVarName . value' <$> genVarName' + (body, x) <- makeOutputFormatFromType' t + body <- (`mapFormatTreeM` body) $ \case + Exp e -> return $ Exp (At e i) + Loop i (Len n) body -> return $ Loop i (Len (At n i)) body + format -> return format + return (Seq [Exp (Len (Var x)), Loop i (Len (Var x)) body], x) + t -> throwSemanticError $ "cannot read input of type: " ++ formatType t + +makeOutputFormatFromType :: (MonadAlpha m, MonadError Error m) => Type -> m (FormatTree, Either String [String]) +makeOutputFormatFromType = \case + TupleTy ts -> do + outputs <- mapM makeOutputFormatFromType' ts + return (Seq (map fst outputs), Right (map snd outputs)) + t -> second Left <$> makeOutputFormatFromType' t + +makeIOFormatFromType :: (MonadAlpha m, MonadError Error m) => [Type] -> Type -> m IOFormat +makeIOFormatFromType ts ret = do + inputs <- mapM makeInputFormatFromType ts + (outputTree, outputVariables) <- makeOutputFormatFromType ret + return $ + IOFormat + { inputTree = Seq (map fst inputs), + inputVariables = map snd inputs, + outputVariables = outputVariables, + outputTree = outputTree + } + +run :: (MonadAlpha m, MonadError Error m) => Program -> m IOFormat +run prog = wrapError' "Jikka.RestrictedPython.Convert.DefaultMain" $ do + (_, args, ret, _) <- lookupSolve prog + makeIOFormatFromType (map snd args) ret diff --git a/src/Jikka/RestrictedPython/Convert/ParseMain.hs b/src/Jikka/RestrictedPython/Convert/ParseMain.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Convert/ParseMain.hs @@ -0,0 +1,225 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PatternSynonyms #-} + +-- | +-- Module : Jikka.RestrictedPython.Convert.ParseMain +-- Description : analyze @main@ function into input formats. / @main@ 関数を分析して入力フォーマットを得ます。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.RestrictedPython.Convert.ParseMain + ( run, + ) +where + +import Control.Arrow +import Data.Maybe +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Common.IOFormat +import Jikka.RestrictedPython.Format (formatExpr, formatTarget) +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Util + +type MainFunction = (Maybe Loc, [(VarName', Type)], Type, [Statement]) + +splitMain :: Program -> (Maybe MainFunction, Program) +splitMain = \case + [] -> (Nothing, []) + ToplevelFunctionDef (WithLoc' loc (VarName "main")) args ret body : stmts -> (Just (loc, args, ret, body), stmts) + stmt : stmts -> second (stmt :) $ splitMain stmts + +checkMainType :: MonadError Error m => MainFunction -> m () +checkMainType (loc, args, ret, _) = wrapAt' loc $ case args of + _ : _ -> throwTypeError "main function must not take arguments" + [] -> case ret of + VarTy _ -> return () + NoneTy -> return () + _ -> throwTypeError "main function must return None" + +pattern CallBuiltin b args <- WithLoc' _ (Call (WithLoc' _ (Constant (ConstBuiltin b))) args) + +pattern CallMethod e a args <- WithLoc' _ (Call (WithLoc' _ (Attribute e a)) args) + +pattern IntInput <- + CallBuiltin (BuiltinInt _) [CallBuiltin BuiltinInput []] + +pattern MapIntInputSplit <- + CallBuiltin + (BuiltinMap [_] _) + [ WithLoc' _ (Constant (ConstBuiltin (BuiltinInt _))), + CallMethod + (CallBuiltin BuiltinInput []) + (WithLoc' _ BuiltinSplit) + [] + ] + +pattern ListMapIntInputSplit <- + CallBuiltin + (BuiltinList _) + [ CallBuiltin + (BuiltinMap [_] _) + [ WithLoc' _ (Constant (ConstBuiltin (BuiltinInt _))), + CallMethod + (CallBuiltin BuiltinInput []) + (WithLoc' _ BuiltinSplit) + [] + ] + ] + +pattern ListRange n <- + CallBuiltin + (BuiltinList _) + [CallBuiltin BuiltinRange1 [WithLoc' _ (Name (WithLoc' _ n))]] + +parseAnnAssign :: (MonadAlpha m, MonadError Error m) => Target' -> Type -> Expr' -> [Statement] -> m (FormatTree, Maybe ([String], Either String [String]), [Statement]) +parseAnnAssign x _ e cont = do + let subscriptTrg x = case value' x of + NameTrg x -> return (unVarName (value' x), []) + SubscriptTrg x (WithLoc' _ (Name i)) -> second (++ [unVarName (value' i)]) <$> subscriptTrg x + _ -> throwSemanticErrorAt' (loc' x) $ "name target or subscript target is expected, but got: " ++ formatTarget x + let subscriptTupleTrg x = case value' x of + TupleTrg xs -> mapM subscriptTrg xs + _ -> throwSemanticErrorAt' (loc' x) $ "tuple target is expected, but got: " ++ formatTarget x + let nameTrg x = case value' x of + NameTrg x -> return $ unVarName (value' x) + _ -> throwSemanticErrorAt' (loc' x) $ "name target is expected, but got: " ++ formatTarget x + let nameOrTupleTrg x = case value' x of + NameTrg x -> return . Left $ unVarName (value' x) + TupleTrg xs -> Right <$> mapM nameTrg xs + _ -> throwSemanticErrorAt' (loc' x) $ "name target or tuple target is expected, but got: " ++ formatTarget x + let nameExpr e = case value' e of + Name x -> return $ unVarName (value' x) + _ -> throwSemanticErrorAt' (loc' e) $ "variable is expected, but got: " ++ formatExpr e + case e of + -- int(input()) + IntInput -> do + (x, indices) <- subscriptTrg x + return (Seq [packSubscriptedVar' x indices, Newline], Nothing, cont) + -- map(int, input().split()) + MapIntInputSplit -> do + outputs <- subscriptTupleTrg x + return (Seq (map (uncurry packSubscriptedVar') outputs ++ [Newline]), Nothing, cont) + -- list(map(int, input().split())) + ListMapIntInputSplit -> do + (x, indices) <- subscriptTrg x + case cont of + Assert (WithLoc' _ (Compare (CallBuiltin (BuiltinLen _) [WithLoc' _ (Name x')]) (CmpOp' Eq' _) n)) : cont | unVarName (value' x') == x -> do + i <- unVarName . value' <$> genVarName' + n <- nameExpr n + return (Seq [Loop i (Var n) (Exp (At (packSubscriptedVar x indices) i)), Newline], Nothing, cont) + _ -> throwSemanticErrorAt' (loc' e) "after `xs = list(map(int, input().split()))', we need to write `assert len(xs) == n`" + -- list(range(n)) + ListRange n -> do + let isListRange = \case + AnnAssign _ _ (ListRange n') | n' == n -> True + _ -> False + cont <- return $ dropWhile isListRange cont + case cont of + For _ (CallBuiltin BuiltinRange1 [WithLoc' _ (Name n')]) _ : _ | value' n' == n -> return (Seq [], Nothing, cont) -- TODO: add more strict checks + _ -> throwSemanticErrorAt' (loc' e) "after some repetition of `xs = list(range(n))', we need to write `for i in range(n):`" + -- solve(...) + WithLoc' _ (Call (WithLoc' _ (Name (WithLoc' _ (VarName "solve")))) args) -> do + inputs <- mapM nameExpr args + output <- nameOrTupleTrg x + return (Seq [], Just (inputs, output), cont) + _ -> throwSemanticErrorAt' (loc' e) "assignments in main function must be `x = int(input())', `x, y, z = map(int, input().split())', `xs = list(map(int, input().split()))', `xs = list(range(n))' or `x, y, z = solve(a, b, c)'" + +parseFor :: MonadError Error m => ([Statement] -> m (FormatTree, Maybe ([String], Either String [String]), FormatTree)) -> Target' -> Expr' -> [Statement] -> m (FormatTree, FormatTree) +parseFor go x e body = do + x <- case value' x of + NameTrg x -> return x + _ -> throwSemanticErrorAt' (loc' x) $ "for loops in main function must use `range' like `for i in range(n): ...'" ++ formatTarget x + n <- case e of + CallBuiltin BuiltinRange1 [n] -> return n + _ -> throwSemanticErrorAt' (loc' e) $ "for loops in main function must use `range' like `for i in range(n): ...': " ++ formatExpr e + n <- case value' n of + Name n -> return $ Right (n, 0) + BinOp (WithLoc' _ (Name n)) Add (WithLoc' _ (Constant (ConstInt k))) -> return $ Right (n, k) + BinOp (WithLoc' _ (Name n)) Sub (WithLoc' _ (Constant (ConstInt k))) -> return $ Right (n, - k) + Call (WithLoc' _ (Constant (ConstBuiltin (BuiltinLen _)))) [WithLoc' _ (Name xs)] -> return $ Left xs + _ -> throwSemanticErrorAt' (loc' n) $ "for loops in main function must use `range(x)', `range(x + k)', `range(x - k)', `range(len(xs))`: " ++ formatExpr n + n <- return $ case n of + Right (n, k) -> + let n' = Var (unVarName (value' n)) + in if k == 0 then n' else Plus n' k + Left xs -> Len (Var (unVarName (value' xs))) + (input, solve, output) <- go body + when (isJust solve) $ do + throwSemanticError "cannot call `solve(...)' in for loop" + let x' = unVarName (value' x) + return (Loop x' n input, Loop x' n output) + +parseExprStatement :: (MonadAlpha m, MonadError Error m) => Expr' -> m FormatTree +parseExprStatement e = do + let subscriptExpr e = case value' e of + Name x -> return (unVarName (value' x), []) + Subscript e (WithLoc' _ (Name i)) -> second (++ [unVarName (value' i)]) <$> subscriptExpr e + _ -> throwSemanticErrorAt' (loc' e) $ "subscripted variable is expected, but got: " ++ formatExpr e + let starredExpr e = do + (e, starred) <- return $ case value' e of + Starred e -> (e, True) + _ -> (e, False) + (x, indices) <- subscriptExpr e + return (x, indices, starred) + let pack (x, indices, starred) + | not starred = return $ packSubscriptedVar' x indices + | otherwise = do + let xs = packSubscriptedVar x indices + i <- unVarName . value' <$> genVarName' + return $ Loop i (Len xs) (packSubscriptedVar' x (indices ++ [i])) + case e of + CallBuiltin (BuiltinPrint _) args -> do + outputs <- mapM starredExpr args + outputs <- mapM pack outputs + return $ Seq (outputs ++ [Newline]) + _ -> throwSemanticErrorAt' (loc' e) "only `print(...)' is allowed for expr statements in main function" + +parseMain :: (MonadAlpha m, MonadError Error m) => MainFunction -> m IOFormat +parseMain (loc, _, _, body) = wrapAt' loc $ pack =<< go [] body + where + pack :: MonadError Error m => (FormatTree, Maybe ([String], Either String [String]), FormatTree) -> m IOFormat + pack (_, Nothing, _) = throwSemanticError "main function must call solve function" + pack (inputTree, Just (inputVariables, outputVariables), outputTree) = + return $ + IOFormat + { inputTree = inputTree, + inputVariables = inputVariables, + outputVariables = outputVariables, + outputTree = outputTree + } + go :: (MonadAlpha m, MonadError Error m) => [(FormatTree, Maybe ([String], Either String [String]), FormatTree)] -> [Statement] -> m (FormatTree, Maybe ([String], Either String [String]), FormatTree) + go formats = \case + Return _ : _ -> throwSemanticError "return statement is not allowd in main function" + AugAssign _ _ _ : _ -> throwSemanticError "augumented assignment statement is not allowd in main function" + AnnAssign x t e : cont -> do + (inputs, solve, cont) <- parseAnnAssign x t e cont + go (formats ++ [(inputs, solve, Seq [])]) cont + For x e body : cont -> do + (inputs, outputs) <- parseFor (go []) x e body + go (formats ++ [(inputs, Nothing, outputs)]) cont + If _ _ _ : _ -> throwSemanticError "if statement is not allowd in main function" + Assert _ : _ -> throwSemanticError "assert statement is allowd only after `xs = list(map(int, input().split()))` in main function" + Expr' e : cont -> do + output <- parseExprStatement e + go (formats ++ [(Seq [], Nothing, output)]) cont + [] -> do + let input = Seq (map (\(x, _, _) -> x) formats) + let outputs = Seq (map (\(_, _, z) -> z) formats) + solve <- case mapMaybe (\(_, y, _) -> y) formats of + [] -> return Nothing + [solve] -> return $ Just solve + _ -> throwSemanticError "cannot call solve function twice" + return (input, solve, outputs) + +run :: (MonadAlpha m, MonadError Error m) => Program -> m (Maybe IOFormat, Program) +run prog = wrapError' "Jikka.RestrictedPython.Convert.ParseMain" $ do + (main, prog) <- return $ splitMain prog + main <- forM main $ \main -> do + checkMainType main + main <- parseMain main + return $ normalizeIOFormat main + return (main, prog) diff --git a/src/Jikka/RestrictedPython/Convert/RemoveUnbalancedIf.hs b/src/Jikka/RestrictedPython/Convert/RemoveUnbalancedIf.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Convert/RemoveUnbalancedIf.hs @@ -0,0 +1,46 @@ +-- | +-- Module : Jikka.RestrictedPython.Convert.RemoveUnbalancedIf +-- Description : converts and removes if-statements whose either branch has return-statements and the other branch doesn't have return-statements. / その一方の分岐は return 文を持ちもう一方の分岐は return 文を持たないような if 文を変形し削除します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.RestrictedPython.Convert.RemoveUnbalancedIf + ( run, + ) +where + +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Util + +runStatements :: [Statement] -> [Statement] +runStatements [] = [] +runStatements (stmt : stmts) = case stmt of + If e body1 body2 -> case (any doesAlwaysReturn body1, any doesAlwaysReturn body2) of + (True, False) -> [If e body1 (body2 ++ runStatements stmts)] + (False, True) -> [If e (body1 ++ runStatements stmts) body2] + _ -> stmt : runStatements stmts + _ -> stmt : runStatements stmts + +-- | `run` removes if-statements that one branch always returns and the other branch doesn't. +-- +-- For example, the following +-- +-- > if True: +-- > return 0 +-- > else: +-- > a = 0 +-- > b = 1 +-- > return 2 +-- +-- is converted to +-- +-- > if True: +-- > return 0 +-- > else: +-- > a = 0 +-- > b = 1 +-- > return 2 +run :: Program -> Program +run = mapStatements runStatements diff --git a/src/Jikka/RestrictedPython/Convert/RemoveUnreachable.hs b/src/Jikka/RestrictedPython/Convert/RemoveUnreachable.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Convert/RemoveUnreachable.hs @@ -0,0 +1,44 @@ +-- | +-- Module : Jikka.RestrictedPython.Convert.RemoveUnreachable +-- Description : removes unreachable statements. / 到達不能な文を削除します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.RestrictedPython.Convert.RemoveUnreachable + ( run, + ) +where + +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Util + +runStatements :: [Statement] -> [Statement] +runStatements stmts = case break doesAlwaysReturn stmts of + (stmts, []) -> stmts + (stmts, stmt : _) -> stmts ++ [stmt] + +-- | `run` removes unreachable statements after return-statements. +-- +-- For example, the following +-- +-- > a = 0 +-- > if True: +-- > b = 0 +-- > return b +-- > b += 1 +-- > else: +-- > return 1 +-- > a += 1 +-- +-- is converted to +-- +-- > a = 0 +-- > if True: +-- > b = 0 +-- > return b +-- > else: +-- > return 1 +run :: Program -> Program +run = mapStatements runStatements diff --git a/src/Jikka/RestrictedPython/Convert/ResolveBuiltin.hs b/src/Jikka/RestrictedPython/Convert/ResolveBuiltin.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Convert/ResolveBuiltin.hs @@ -0,0 +1,43 @@ +{-# LANGUAGE FlexibleContexts #-} + +-- | +-- Module : Jikka.RestrictedPython.Convert.ResolveBuiltin +-- Description : resolves names of builtin functions using information of arity. / arity の情報を使いながら組み込み関数を名前解決します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.RestrictedPython.Convert.ResolveBuiltin + ( run, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.RestrictedPython.Language.Builtin +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Lint +import Jikka.RestrictedPython.Language.Util + +runExpr :: (MonadAlpha m, MonadError Error m) => Expr' -> m Expr' +runExpr = mapSubExprM go + where + go :: (MonadAlpha m, MonadError Error m) => Expr' -> m Expr' + go e = case value' e of + Name x -> resolveUniqueBuiltin x + Call (WithLoc' _ (Name f)) args -> WithLoc' (loc' e) <$> (Call <$> resolveBuiltin f (length args) <*> pure args) + Attribute e' a -> WithLoc' (loc' e) <$> resolveAttribute e' a + _ -> return e + +-- | `run` resolves types of polymorphic builtin functions. +-- This assumes there are no assignments to builtin functions, i.e. `doesntHaveAssignmentToBuiltin`. +-- +-- For example, the @max@ of @max(xs)@ has a type \(\mathbf{list}(\alpha) \to \alpha\) but the @max@ of @max(x, y, z)@ has a type \(\alpha \times \alpha \times \alpha \to \alpha\). +-- So this function converts @Var "max"@ to @BuiltinMax1 t@, @BuiltinMax t 2@, @BuiltinMax t 3@, etc.. +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.RestrictedPython.Convert.ResolveBuiltin" $ do + ensureDoesntHaveAssignmentToBuiltin prog + prog <- mapExprM runExpr prog + ensureDoesntHaveNonResolvedBuiltin prog + return prog diff --git a/src/Jikka/RestrictedPython/Convert/SplitLoops.hs b/src/Jikka/RestrictedPython/Convert/SplitLoops.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Convert/SplitLoops.hs @@ -0,0 +1,74 @@ +{-# LANGUAGE FlexibleContexts #-} + +-- | +-- Module : Jikka.RestrictedPython.Convert.TypeInfer +-- Description : split a for-loop into many small for-loops based on the dependency graph of variables and assignments. / 変数と代入の依存関係グラフに基づいて、for ループを複数の小さな for ループに分割します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.RestrictedPython.Convert.SplitLoops + ( run, + run', + runForLoop, + ) +where + +import Data.List (partition) +import Jikka.Common.Alpha +import Jikka.Common.Error +import qualified Jikka.RestrictedPython.Convert.Alpha as Alpha (run) +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Lint +import Jikka.RestrictedPython.Language.Util +import Jikka.RestrictedPython.Language.VariableAnalysis + +-- | `runForLoop` splits a for-loop to many for-loops as possible. +-- This assumes that `doesntHaveSubscriptionInLoopCounters`, `doesntHaveAssignmentToLoopCounters`, and `doesntHaveAssignmentToLoopIterators` hold. +-- +-- This function analyzes read-variables and write-variables in statements, and split statements into connected components. +runForLoop :: Target' -> Expr' -> [Statement] -> [Statement] +runForLoop x iter body = + let connected (_, (r, w)) (_, (r', w')) = haveWriteReadIntersection w r' || haveWriteReadIntersection w' r + go result [] = reverse result + go result (stmt : stmts) = + let (same, diff) = partition (connected stmt) stmts + in go (For x iter (map fst (stmt : same)) : result) diff + body' = map (\stmt -> (stmt, analyzeStatementMax stmt)) body + in go [] body' + +-- | `run'` splits for-loops into many small for-loops as possible. +-- This assumes that `doesntHaveSubscriptionInLoopCounters`, `doesntHaveAssignmentToLoopCounters`, and `doesntHaveAssignmentToLoopIterators` hold. +-- This may introduce name conflicts. +-- +-- For example, the following +-- +-- > a = 0 +-- > b = 0 +-- > for i in range(10): +-- > c = b +-- > a += i +-- > b += c +-- +-- is split to +-- +-- > a = 0 +-- > b = 0 +-- > for i in range(10): +-- > c = b +-- > b += c +-- > for i in range(10): +-- > a += i +run' :: Program -> Program +run' = mapLargeStatement (\e pred1 pred2 -> [If e pred1 pred2]) runForLoop + +-- | `run` does alpha conversion, check assumptions, and `run'`. +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.RestrictedPython.Convert.SplitLoops" $ do + prog <- Alpha.run prog + ensureDoesntHaveSubscriptionInLoopCounters prog + ensureDoesntHaveAssignmentToLoopCounters prog + ensureDoesntHaveAssignmentToLoopIterators prog + prog <- return $ run' prog + Alpha.run prog diff --git a/src/Jikka/RestrictedPython/Convert/ToCore.hs b/src/Jikka/RestrictedPython/Convert/ToCore.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Convert/ToCore.hs @@ -0,0 +1,465 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TupleSections #-} + +-- | +-- Module : Jikka.RestrictedPython.Convert.ToCore +-- Description : converts programs of our restricted Python to programs of core language. / 制限された Python のプログラムを core 言語のプログラムに変換します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.RestrictedPython.Convert.ToCore + ( run, + runForStatement, + runIfStatement, + ) +where + +import Control.Monad.State.Strict +import Data.List (intersect, union) +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Common.Location +import qualified Jikka.Core.Language.BuiltinPatterns as Y +import qualified Jikka.Core.Language.Expr as Y +import qualified Jikka.Core.Language.Util as Y +import qualified Jikka.RestrictedPython.Language.Expr as X +import qualified Jikka.RestrictedPython.Language.Lint as X +import qualified Jikka.RestrictedPython.Language.Util as X +import qualified Jikka.RestrictedPython.Language.VariableAnalysis as X + +type Env = [X.VarName] + +defineVar :: MonadState Env m => X.VarName -> m () +defineVar x = modify' (x :) + +isDefinedVar :: MonadState Env m => X.VarName -> m Bool +isDefinedVar x = gets (x `elem`) + +withScope :: MonadState Env m => m a -> m a +withScope f = do + env <- get + x <- f + put env + return x + +runVarName :: X.VarName' -> Y.VarName +runVarName (X.WithLoc' _ (X.VarName x)) = Y.VarName x + +runType :: MonadError Error m => X.Type -> m Y.Type +runType = \case + X.VarTy (X.TypeName x) -> return $ Y.VarTy (Y.TypeName x) + X.IntTy -> return Y.IntTy + X.BoolTy -> return Y.BoolTy + X.ListTy t -> Y.ListTy <$> runType t + X.TupleTy ts -> Y.TupleTy <$> mapM runType ts + X.CallableTy args ret -> Y.curryFunTy <$> mapM runType args <*> runType ret + X.StringTy -> throwSemanticError "cannot use `str' type out of main function" + X.SideEffectTy -> throwSemanticError "side-effect type must be used only as expr-statement" -- TODO: check in Jikka.RestrictedPython.Language.Lint + +runConstant :: MonadError Error m => X.Constant -> m Y.Expr +runConstant = \case + X.ConstNone -> return $ Y.Tuple' [] + X.ConstInt n -> return $ Y.Lit (Y.LitInt n) + X.ConstBool p -> return $ Y.Lit (Y.LitBool p) + X.ConstBuiltin builtin -> runBuiltin builtin + +runBuiltin :: MonadError Error m => X.Builtin -> m Y.Expr +runBuiltin builtin = + let f = return . Y.Lit . Y.LitBuiltin + in case builtin of + X.BuiltinAbs -> f Y.Abs + X.BuiltinPow -> f Y.Pow + X.BuiltinModPow -> f Y.ModPow + X.BuiltinDivMod -> return $ Y.Lam2 "a" Y.IntTy "b" Y.IntTy (Y.uncurryApp (Y.Tuple' [Y.IntTy, Y.IntTy]) [Y.FloorDiv' (Y.Var "a") (Y.Var "b"), Y.FloorMod' (Y.Var "a") (Y.Var "b")]) + X.BuiltinCeilDiv -> f Y.CeilDiv + X.BuiltinCeilMod -> f Y.CeilMod + X.BuiltinFloorDiv -> f Y.FloorDiv + X.BuiltinFloorMod -> f Y.FloorMod + X.BuiltinGcd -> f Y.Gcd + X.BuiltinLcm -> f Y.Lcm + X.BuiltinInt t -> case t of + X.IntTy -> return $ Y.Lam "x" Y.IntTy (Y.Var "x") + X.BoolTy -> return $ Y.Lam "p" Y.BoolTy (Y.If' Y.IntTy (Y.Var "p") Y.Lit1 Y.Lit0) + _ -> throwTypeError "the argument of int must be int or bool" + X.BuiltinBool t -> case t of + X.IntTy -> return $ Y.Lam "x" Y.IntTy (Y.If' Y.BoolTy (Y.Equal' Y.IntTy (Y.Var "x") Y.Lit0) Y.LitFalse Y.LitTrue) + X.BoolTy -> return $ Y.Lam "p" Y.BoolTy (Y.Var "p") + X.ListTy t -> do + t <- runType t + return $ Y.Lam "xs" (Y.ListTy t) (Y.If' Y.BoolTy (Y.Equal' (Y.ListTy t) (Y.Var "xs") (Y.Lit (Y.LitNil t))) Y.LitFalse Y.LitTrue) + _ -> throwTypeError "the argument of bool must be bool, int, or list(a)" + X.BuiltinList t -> do + t <- runType t + return $ Y.Lam "xs" (Y.ListTy t) (Y.Var "xs") + X.BuiltinTuple ts -> f . Y.Tuple =<< mapM runType ts + X.BuiltinLen t -> f . Y.Len =<< runType t + X.BuiltinMap ts ret -> case ts of + [] -> Y.Nil' <$> runType ret + _ -> do + ts <- mapM runType ts + ret <- runType ret + let var i = Y.VarName ("xs" ++ show i) + let lam body = Y.Lam "f" (Y.curryFunTy ts ret) (foldr (\(i, t) -> Y.Lam (var i) (Y.ListTy t)) body (zip [0 ..] ts)) + let len = Y.Min1' Y.IntTy (foldr (Y.Cons' Y.IntTy) (Y.Nil' Y.IntTy) (zipWith (\i t -> Y.Len' t (Y.Var (var i))) [0 ..] ts)) + let body = Y.Map' Y.IntTy ret (Y.Lam "i" Y.IntTy (Y.uncurryApp (Y.Var "f") (map (Y.Var . var) [0 .. length ts - 1]))) (Y.Range1' len) + return $ lam body + X.BuiltinSorted t -> f . Y.Sorted =<< runType t + X.BuiltinReversed t -> f . Y.Reversed =<< runType t + X.BuiltinEnumerate t -> do + t <- runType t + let body = Y.Lam "i" Y.IntTy (Y.uncurryApp (Y.Tuple' [Y.IntTy, t]) [Y.Var "i", Y.At' t (Y.Var "xs") (Y.Var "i")]) + return $ Y.Lam "xs" (Y.ListTy t) (Y.Map' (Y.ListTy t) (Y.ListTy (Y.TupleTy [Y.IntTy, t])) body (Y.Range1' (Y.Len' t (Y.Var "xs")))) + X.BuiltinFilter t -> f . Y.Filter =<< runType t + X.BuiltinZip ts -> do + ts <- mapM runType ts + let var i = Y.VarName ("xs" ++ show i) + let lam body = foldr (\(i, t) -> Y.Lam (var i) (Y.ListTy t)) body (zip [0 ..] ts) + let len = Y.Min1' Y.IntTy (foldr (Y.Cons' Y.IntTy) (Y.Nil' Y.IntTy) (zipWith (\i t -> Y.Len' t (Y.Var (var i))) [0 ..] ts)) + let body = Y.Map' Y.IntTy (Y.TupleTy ts) (Y.Lam "i" Y.IntTy (Y.uncurryApp (Y.Tuple' ts) (map (Y.Var . var) [0 .. length ts - 1]))) (Y.Range1' len) + return $ lam body + X.BuiltinAll -> f Y.All + X.BuiltinAny -> f Y.Any + X.BuiltinSum -> f Y.Sum + X.BuiltinProduct -> f Y.Product + X.BuiltinRange1 -> f Y.Range1 + X.BuiltinRange2 -> f Y.Range2 + X.BuiltinRange3 -> f Y.Range1 + X.BuiltinMax1 t -> f . Y.Max1 =<< runType t + X.BuiltinMax t n -> do + when (n < 2) $ do + throwTypeError $ "max expected 2 or more arguments, got " ++ show n + t <- runType t + let args = map (\i -> Y.VarName ('x' : show i)) [0 .. n -1] + return $ Y.curryLam (map (,t) args) (foldr1 (Y.Max2' t) (map Y.Var args)) + X.BuiltinMin1 t -> f . Y.Min1 =<< runType t + X.BuiltinMin t n -> do + when (n < 2) $ do + throwTypeError $ "max min 2 or more arguments, got " ++ show n + t <- runType t + let args = map (\i -> Y.VarName ('x' : show i)) [0 .. n -1] + return $ Y.curryLam (map (,t) args) (foldr1 (Y.Min2' t) (map Y.Var args)) + X.BuiltinArgMax t -> f . Y.ArgMax =<< runType t + X.BuiltinArgMin t -> f . Y.ArgMin =<< runType t + X.BuiltinFact -> f Y.Fact + X.BuiltinChoose -> f Y.Choose + X.BuiltinPermute -> f Y.Permute + X.BuiltinMultiChoose -> f Y.MultiChoose + X.BuiltinModInv -> f Y.ModInv + X.BuiltinInput -> throwSemanticError "cannot use `input' out of main function" + X.BuiltinPrint _ -> throwSemanticError "cannot use `print' out of main function" + +runAttribute :: MonadError Error m => X.Attribute' -> m Y.Expr +runAttribute a = wrapAt' (loc' a) $ do + case value' a of + X.UnresolvedAttribute a -> throwInternalError $ "unresolved attribute: " ++ X.unAttributeName a + X.BuiltinCount t -> do + t <- runType t + return $ Y.Lam2 "xs" (Y.ListTy t) "x" t (Y.Len' t (Y.Filter' t (Y.Lam "y" t (Y.Equal' t (Y.Var "x") (Y.Var "y"))) (Y.Var "xs"))) + X.BuiltinIndex t -> do + t <- runType t + return $ Y.Lam2 "xs" (Y.ListTy t) "x" t (Y.Min1' Y.IntTy (Y.Filter' Y.IntTy (Y.Lam "i" Y.IntTy (Y.Equal' t (Y.At' t (Y.Var "xs") (Y.Var "i")) (Y.Var "x"))) (Y.Range1' (Y.Len' t (Y.Var "xs"))))) + X.BuiltinCopy t -> do + t <- runType t + return $ Y.Lam "x" t (Y.Var "x") + X.BuiltinAppend _ -> throwSemanticError "cannot use `append' out of expr-statements" + X.BuiltinSplit -> throwSemanticError "cannot use `split' out of main function" + +runBoolOp :: X.BoolOp -> Y.Builtin +runBoolOp = \case + X.And -> Y.And + X.Or -> Y.Or + X.Implies -> Y.Implies + +runUnaryOp :: X.UnaryOp -> Y.Expr +runUnaryOp = + let f = Y.Lit . Y.LitBuiltin + in \case + X.Invert -> f Y.BitNot + X.Not -> f Y.Not + X.UAdd -> Y.Lam "x" Y.IntTy (Y.Var "x") + X.USub -> f Y.Negate + +runOperator :: MonadError Error m => X.Operator -> m Y.Builtin +runOperator = \case + X.Add -> return Y.Plus + X.Sub -> return Y.Minus + X.Mult -> return Y.Mult + X.MatMult -> throwSemanticError "matmul operator ('@') is not supported" + X.Div -> throwSemanticError "floatdiv operator ('/') is not supported" + X.FloorDiv -> return Y.FloorDiv + X.FloorMod -> return Y.FloorMod + X.CeilDiv -> return Y.CeilDiv + X.CeilMod -> return Y.CeilMod + X.Pow -> return Y.Pow + X.BitLShift -> return Y.BitLeftShift + X.BitRShift -> return Y.BitRightShift + X.BitOr -> return Y.BitOr + X.BitXor -> return Y.BitXor + X.BitAnd -> return Y.BitAnd + X.Max -> return $ Y.Max2 Y.IntTy + X.Min -> return $ Y.Min2 Y.IntTy + +runCmpOp :: MonadError Error m => X.CmpOp' -> m Y.Expr +runCmpOp (X.CmpOp' op t) = do + t <- runType t + let f = Y.Lit . Y.LitBuiltin + return $ case op of + X.Lt -> f $ Y.LessThan t + X.LtE -> f $ Y.LessEqual t + X.Gt -> f $ Y.GreaterThan t + X.GtE -> f $ Y.GreaterEqual t + X.Eq' -> f $ Y.Equal t + X.NotEq -> f $ Y.NotEqual t + X.Is -> f $ Y.Equal t + X.IsNot -> f $ Y.NotEqual t + X.In -> f $ Y.Elem t + X.NotIn -> Y.curryLam [("x", t), ("xs", Y.ListTy t)] (Y.Not' (Y.Elem' t (Y.Var "x") (Y.Var "xs"))) + +runTargetExpr :: (MonadAlpha m, MonadError Error m) => X.Target' -> m Y.Expr +runTargetExpr (WithLoc' _ x) = case x of + X.SubscriptTrg x e -> Y.At' <$> Y.genType <*> runTargetExpr x <*> runExpr e + X.NameTrg x -> return $ Y.Var (runVarName x) + X.TupleTrg xs -> Y.uncurryApp <$> (Y.Tuple' <$> replicateM (length xs) Y.genType) <*> mapM runTargetExpr xs + +runAssign :: (MonadAlpha m, MonadError Error m) => X.Target' -> Y.Expr -> m Y.Expr -> m Y.Expr +runAssign (WithLoc' _ x) e cont = case x of + X.SubscriptTrg x index -> join $ runAssign x <$> (Y.SetAt' <$> Y.genType <*> runTargetExpr x <*> runExpr index <*> pure e) <*> pure cont + X.NameTrg x -> Y.Let (runVarName x) <$> Y.genType <*> pure e <*> cont + X.TupleTrg xs -> do + y <- Y.genVarName' + ts <- replicateM (length xs) Y.genType + cont <- join $ foldM (\cont (i, x) -> return $ runAssign x (Y.Proj' ts i (Y.Var y)) cont) cont (zip [0 ..] xs) + return $ Y.Let y (Y.TupleTy ts) e cont + +runListComp :: (MonadAlpha m, MonadError Error m) => X.Expr' -> X.Comprehension -> m Y.Expr +runListComp e (X.Comprehension x iter pred) = do + iter <- runExpr iter + y <- Y.genVarName' + t1 <- Y.genType + iter <- case pred of + Nothing -> return iter + Just pred -> Y.Filter' t1 <$> (Y.Lam y t1 <$> runAssign x (Y.Var y) (runExpr pred)) <*> pure iter + t2 <- Y.genType + e <- runExpr e + Y.Map' t1 t2 <$> (Y.Lam y t1 <$> runAssign x (Y.Var y) (pure e)) <*> pure iter + +runExpr :: (MonadAlpha m, MonadError Error m) => X.Expr' -> m Y.Expr +runExpr e0 = wrapAt' (loc' e0) $ case value' e0 of + X.BoolOp e1 op e2 -> Y.AppBuiltin2 (runBoolOp op) <$> runExpr e1 <*> runExpr e2 + X.BinOp e1 op e2 -> Y.AppBuiltin2 <$> runOperator op <*> runExpr e1 <*> runExpr e2 + X.UnaryOp op e -> Y.App (runUnaryOp op) <$> runExpr e + X.Lambda args body -> Y.curryLam <$> mapM (\(x, t) -> (runVarName x,) <$> runType t) args <*> runExpr body + X.IfExp e1 e2 e3 -> do + e1 <- runExpr e1 + e2 <- runExpr e2 + e3 <- runExpr e3 + t <- Y.genType + return $ Y.If' t e1 e2 e3 + X.ListComp x comp -> runListComp x comp + X.Compare e1 op e2 -> Y.App2 <$> runCmpOp op <*> runExpr e1 <*> runExpr e2 + X.Call f args -> Y.uncurryApp <$> runExpr f <*> mapM runExpr args + X.Constant const -> runConstant const + X.Attribute e a -> do + e <- runExpr e + a <- runAttribute a + return $ Y.App a e + X.Subscript e1 e2 -> Y.AppBuiltin2 <$> (Y.At <$> Y.genType) <*> runExpr e1 <*> runExpr e2 + X.Starred e -> throwSemanticErrorAt' (loc' e) "cannot use starred expr" + X.Name x -> return $ Y.Var (runVarName x) + X.List t es -> do + t <- runType t + foldr (Y.Cons' t) (Y.Lit (Y.LitNil t)) <$> mapM runExpr es + X.Tuple es -> Y.uncurryApp <$> (Y.Tuple' <$> mapM (const Y.genType) es) <*> mapM runExpr es + X.SubscriptSlice e from to step -> do + e <- runExpr e + from <- traverse runExpr from + to <- traverse runExpr to + step <- traverse runExpr step + i <- Y.genVarName' + t <- Y.genType + let mapAt = return . Y.Map' Y.IntTy t (Y.Lam i t (Y.At' t e (Y.Var i))) + case (from, to, step) of + (Nothing, Nothing, Nothing) -> return e + (Nothing, Just to, Nothing) -> mapAt (Y.Range1' to) + (Just from, Nothing, Nothing) -> mapAt (Y.Range2' from (Y.Len' t e)) + (Just from, Just to, Nothing) -> mapAt (Y.Range2' from to) + (Nothing, Nothing, Just step) -> mapAt (Y.Range3' Y.Lit0 (Y.Len' t e) step) + (Nothing, Just to, Just step) -> mapAt (Y.Range3' Y.Lit0 to step) + (Just from, Nothing, Just step) -> mapAt (Y.Range3' from (Y.Len' t e) step) + (Just from, Just to, Just step) -> mapAt (Y.Range3' from to step) + +-- | `runForStatement` converts for-loops to `foldl`. +-- For example, this converts the following: +-- +-- > # a, b are defined +-- > for _ in range(n): +-- > c = a + b +-- > a = b +-- > b = c +-- > ... +-- +-- to: +-- +-- > let (a, b) = foldl (fun (a, b) i -> (b, a + b)) (a, b) (range n) +-- > in ... +runForStatement :: (MonadState Env m, MonadAlpha m, MonadError Error m) => X.Target' -> X.Expr' -> [X.Statement] -> [X.Statement] -> [[X.Statement]] -> m Y.Expr +runForStatement x iter body cont conts = do + tx <- Y.genType + iter <- runExpr iter + x' <- Y.genVarName' + z <- Y.genVarName' + let (_, X.WriteList w) = X.analyzeStatementsMax body + ys <- filterM isDefinedVar w + ts <- replicateM (length ys) Y.genType + let init = Y.uncurryApp (Y.Tuple' ts) (map (Y.Var . runVarName . withoutLoc) ys) + let write cont = foldr (\(i, y, t) -> Y.Let (runVarName $ X.WithLoc' Nothing y) t (Y.Proj' ts i (Y.Var z))) cont (zip3 [0 ..] ys ts) + body <- runAssign x (Y.Var x') $ do + runStatements (body ++ [X.Return (withoutLoc (X.Tuple (map (withoutLoc . X.Name . withoutLoc) ys)))]) (cont : conts) + let loop init = Y.Foldl' tx (Y.TupleTy ts) (Y.Lam2 z (Y.TupleTy ts) x' tx (write body)) init iter + cont <- runStatements cont conts + return $ Y.Let z (Y.TupleTy ts) (loop init) (write cont) + +-- | `runIfStatement` converts if-loops to if-exprs. +-- +-- > # a, b are defined +-- > if True: +-- > a = 0 +-- > b = 1 +-- > c = 3 +-- > else: +-- > a = 1 +-- > c = 10 +-- > ... +-- +-- to: +-- +-- > let (a, c) = if true then (0, 3) else (1, 10) +-- > in ... +runIfStatement :: (MonadState Env m, MonadAlpha m, MonadError Error m) => X.Expr' -> [X.Statement] -> [X.Statement] -> [X.Statement] -> [[X.Statement]] -> m Y.Expr +runIfStatement e body1 body2 cont conts = do + e <- runExpr e + t <- Y.genType + case (any X.doesAlwaysReturn body1, any X.doesAlwaysReturn body2) of + (False, False) -> do + let (_, X.WriteList w1) = X.analyzeStatementsMin body1 + let (_, X.WriteList w2) = X.analyzeStatementsMin body2 + let (X.ReadList r, _) = X.analyzeStatementsMax (concat (cont : conts)) + let w = (r `intersect` w1) `union` (r `intersect` w2) + let read = withoutLoc (X.Tuple (map (withoutLoc . X.Name . withoutLoc) w)) + ts <- replicateM (length w) Y.genType + z <- Y.genVarName' + let write value cont = Y.Let z (Y.TupleTy ts) value (foldr (\(i, y, t) -> Y.Let (runVarName (withoutLoc y)) t (Y.Proj' ts i (Y.Var z))) cont (zip3 [0 ..] w ts)) + body1 <- runStatements (body1 ++ [X.Return read]) (cont : conts) + body2 <- runStatements (body2 ++ [X.Return read]) (cont : conts) + cont <- runStatements cont conts + return $ write (Y.If' t e body1 body2) cont + (False, True) -> Y.If' t e <$> runStatements (body1 ++ cont) conts <*> runStatements body2 [] + (True, False) -> Y.If' t e <$> runStatements body1 [] <*> runStatements (body2 ++ cont) conts + (True, True) -> Y.If' t e <$> runStatements body1 [] <*> runStatements body2 [] + +runStatements :: (MonadState Env m, MonadAlpha m, MonadError Error m) => [X.Statement] -> [[X.Statement]] -> m Y.Expr +runStatements [] _ = throwSemanticError "function may not return" +runStatements (stmt : stmts) cont = case stmt of + X.Return e -> runExpr e + X.AugAssign x op e -> do + y <- runTargetExpr x + op <- Y.Lit . Y.LitBuiltin <$> runOperator op + e <- runExpr e + runAssign x (Y.App2 op y e) $ do + runStatements stmts cont + X.AnnAssign x _ e -> do + e <- runExpr e + runAssign x e $ do + withScope $ do + mapM_ defineVar (X.targetVars x) + runStatements stmts cont + X.For x iter body -> runForStatement x iter body stmts cont + X.If e body1 body2 -> runIfStatement e body1 body2 stmts cont + X.Assert _ -> runStatements stmts cont + X.Append loc t x e -> do + case X.exprToTarget x of + Nothing -> throwSemanticErrorAt' loc "invalid `append` method" + Just x -> do + t <- runType t + y <- runTargetExpr x + e <- runExpr e + runAssign x (Y.Snoc' t y e) $ do + runStatements stmts cont + X.Expr' e -> throwSemanticErrorAt' (loc' e) "invalid expr-statement" + +runToplevelStatements :: (MonadState Env m, MonadAlpha m, MonadError Error m) => [X.ToplevelStatement] -> m Y.ToplevelExpr +runToplevelStatements [] = return $ Y.ResultExpr (Y.Var "solve") +runToplevelStatements (stmt : stmts) = case stmt of + X.ToplevelAnnAssign x t e -> do + e <- runExpr e + defineVar (X.value' x) + cont <- runToplevelStatements stmts + t <- runType t + return $ Y.ToplevelLet (runVarName x) t e cont + X.ToplevelFunctionDef f args ret body -> do + defineVar (X.value' f) + body <- withScope $ do + mapM_ (defineVar . X.value' . fst) args + runStatements body [] + cont <- runToplevelStatements stmts + args <- mapM (\(x, t) -> (runVarName x,) <$> runType t) args + ret <- runType ret + return $ Y.ToplevelLetRec (runVarName f) args ret body cont + X.ToplevelAssert _ -> runToplevelStatements stmts -- TOOD: use assertions as hints + +-- | `run` converts programs of our restricted Python-like language to programs of our core language. +-- This assumes the follwing conditions: +-- +-- * `X.doesntHaveSubscriptionInLoopCounters` +-- * `X.doesntHaveLeakOfLoopCounters` +-- * `X.doesntHaveAssignmentToLoopCounters` +-- * `X.doesntHaveAssignmentToLoopIterators` +-- * `X.doesntHaveReturnInLoops` +-- * `X.doesntHaveNonTrivialSubscriptedAssignmentInForLoops` +-- +-- For example, this converts the following: +-- +-- > def solve(n): +-- > if n == 0: +-- > return 1 +-- > else: +-- > return n * solve(n - 1) +-- +-- to: +-- +-- > let solve n = +-- > if n == 0 then +-- > 1 +-- > else: +-- > n * solve (n - 1) +-- > in solve +-- +-- Also, this converts the following: +-- +-- > def solve(n): +-- > a = 0 +-- > b = 1 +-- > for _ in range(n): +-- > c = a + b +-- > a = b +-- > b = c +-- > return a +-- +-- to: +-- +-- > let solve n = +-- > fst (foldl (fun (a, b) i -> (b, a + b)) (0, 1) [0 .. n - 1]) +-- > in solve +run :: (MonadAlpha m, MonadError Error m) => X.Program -> m Y.Program +run prog = wrapError' "Jikka.RestrictedPython.Convert.ToCore" $ do + X.ensureDoesntHaveSubscriptionInLoopCounters prog + X.ensureDoesntHaveLeakOfLoopCounters prog + X.ensureDoesntHaveAssignmentToLoopCounters prog + X.ensureDoesntHaveAssignmentToLoopIterators prog + X.ensureDoesntHaveReturnInLoops prog + X.ensureDoesntHaveNonTrivialSubscriptedAssignmentInForLoops prog + evalStateT (runToplevelStatements prog) [] diff --git a/src/Jikka/RestrictedPython/Convert/TypeInfer.hs b/src/Jikka/RestrictedPython/Convert/TypeInfer.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Convert/TypeInfer.hs @@ -0,0 +1,352 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.RestrictedPython.Convert.TypeInfer +-- Description : does type inference. / 型推論を行います。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.RestrictedPython.Convert.TypeInfer + ( run, + + -- * internal types and functions + Equation (..), + formularizeProgram, + sortEquations, + mergeAssertions, + Subst (..), + subst, + solveEquations, + mapTypeProgram, + ) +where + +import Control.Arrow (second) +import Control.Monad.Reader +import Control.Monad.State.Strict +import Control.Monad.Writer.Strict +import qualified Data.Map.Strict as M +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.RestrictedPython.Format (formatType) +import Jikka.RestrictedPython.Language.Builtin +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Util + +data Equation + = TypeEquation Type Type (Maybe Loc) + | TypeAssertion VarName' Type + deriving (Eq, Ord, Show, Read) + +type Eqns = Dual [Equation] + +formularizeType :: MonadWriter Eqns m => Type -> Type -> Maybe Loc -> m () +formularizeType t1 t2 location = tell $ Dual [TypeEquation t1 t2 location] + +formularizeVarName :: MonadWriter Eqns m => VarName' -> Type -> m () +formularizeVarName x t = tell $ Dual [TypeAssertion x t] + +formularizeTarget :: (MonadWriter Eqns m, MonadAlpha m) => Target' -> m Type +formularizeTarget x0 = case value' x0 of + SubscriptTrg f index -> do + t <- genType + tf <- formularizeTarget f + formularizeType tf (ListTy t) (loc' x0) + tindex <- formularizeExpr index + formularizeType tindex IntTy (loc' x0) + return t + NameTrg x -> do + t <- genType + formularizeVarName x t + return t + TupleTrg xs -> do + TupleTy <$> mapM formularizeTarget xs + +formularizeTarget' :: (MonadWriter Eqns m, MonadAlpha m) => Target' -> Type -> m () +formularizeTarget' x0 t = do + t' <- formularizeTarget x0 + formularizeType t t' (loc' x0) + +formularizeExpr :: (MonadWriter Eqns m, MonadAlpha m) => Expr' -> m Type +formularizeExpr e0 = case value' e0 of + BoolOp e1 _ e2 -> do + formularizeExpr' e1 BoolTy + formularizeExpr' e2 BoolTy + return BoolTy + BinOp e1 _ e2 -> do + formularizeExpr' e1 IntTy + formularizeExpr' e2 IntTy + return IntTy + UnaryOp op e -> do + let t' = if op == Not then BoolTy else IntTy + formularizeExpr' e t' + return t' + Lambda args body -> do + mapM_ (uncurry formularizeVarName) args + ret <- genType + formularizeExpr' body ret + return $ CallableTy (map snd args) ret + IfExp e1 e2 e3 -> do + formularizeExpr' e1 BoolTy + t <- formularizeExpr e2 + formularizeExpr' e3 t + return t + ListComp e comp -> do + let Comprehension x iter pred = comp + te <- formularizeExpr e + tx <- formularizeTarget x + formularizeExpr' iter (ListTy tx) + case pred of + Nothing -> return () + Just pred -> formularizeExpr' pred BoolTy + return $ ListTy te + Compare e1 (CmpOp' op t) e2 -> do + formularizeExpr' e1 t + formularizeExpr' e2 (if op == In || op == NotIn then ListTy t else t) + return BoolTy + Call f args -> do + ts <- mapM formularizeExpr args + ret <- genType + formularizeExpr' f (CallableTy ts ret) + return ret + Constant const -> + return $ case const of + ConstNone -> NoneTy + ConstInt _ -> IntTy + ConstBool _ -> BoolTy + ConstBuiltin b -> typeBuiltin b + Attribute e x -> do + let (t1, t2) = typeAttribute (value' x) + formularizeExpr' e t1 + return t2 + Subscript e1 e2 -> do + t <- genType + formularizeExpr' e1 (ListTy t) + formularizeExpr' e2 IntTy + return t + Starred e -> do + t <- genType + formularizeExpr' e (ListTy t) + return t -- because @*xs@ and @y@ has the same type in @[*xs, y]@ + Name x -> do + t <- genType + formularizeVarName x t + return t + List t es -> do + forM_ es $ \e -> do + formularizeExpr' e t + return $ ListTy t + Tuple es -> TupleTy <$> mapM formularizeExpr es + SubscriptSlice e from to step -> do + t' <- genType + formularizeExpr' e (ListTy t') + let formularize = \case + Nothing -> return () + Just e -> formularizeExpr' e IntTy + formularize from + formularize to + formularize step + return (ListTy t') + +formularizeExpr' :: (MonadWriter Eqns m, MonadAlpha m) => Expr' -> Type -> m () +formularizeExpr' e0 t = do + t' <- formularizeExpr e0 + formularizeType t t' (loc' e0) + +formularizeStatement :: (MonadWriter Eqns m, MonadAlpha m) => Type -> Statement -> m () +formularizeStatement ret = \case + Return e -> do + t <- formularizeExpr e + formularizeType t ret (loc' e) + AugAssign x _ e -> do + formularizeTarget' x IntTy + formularizeExpr' e IntTy + AnnAssign x t e -> do + formularizeTarget' x t + formularizeExpr' e t + For x e body -> do + t <- formularizeTarget x + formularizeExpr' e (ListTy t) + mapM_ (formularizeStatement ret) body + If e body1 body2 -> do + formularizeExpr' e BoolTy + mapM_ (formularizeStatement ret) body1 + mapM_ (formularizeStatement ret) body2 + Assert e -> do + formularizeExpr' e BoolTy + Expr' e -> do + formularizeExpr' e SideEffectTy + +formularizeToplevelStatement :: (MonadWriter Eqns m, MonadAlpha m) => ToplevelStatement -> m () +formularizeToplevelStatement = \case + ToplevelAnnAssign x t e -> do + formularizeVarName x t + formularizeExpr' e t + ToplevelFunctionDef f args ret body -> do + mapM_ (uncurry formularizeVarName) args + formularizeVarName f (CallableTy (map snd args) ret) + mapM_ (formularizeStatement ret) body + ToplevelAssert e -> do + formularizeExpr' e BoolTy + +formularizeProgram :: MonadAlpha m => Program -> m [Equation] +formularizeProgram prog = getDual <$> execWriterT (mapM_ formularizeToplevelStatement prog) + +sortEquations :: [Equation] -> ([(Type, Type, Maybe Loc)], [(VarName', Type)]) +sortEquations = go [] [] + where + go eqns' assertions [] = (eqns', assertions) + go eqns' assertions (eqn : eqns) = case eqn of + TypeEquation t1 t2 loc -> go ((t1, t2, loc) : eqns') assertions eqns + TypeAssertion x t -> go eqns' ((x, t) : assertions) eqns + +mergeAssertions :: [(VarName', Type)] -> [(Type, Type, Maybe Loc)] +mergeAssertions = go M.empty [] + where + go _ eqns [] = eqns + go gamma eqns ((x, t) : assertions) = case M.lookup (value' x) gamma of + Nothing -> go (M.insert (value' x) t gamma) eqns assertions + Just t' -> go gamma ((t, t', loc' x) : eqns) assertions + +-- | `Subst` is type substituion. It's a mapping from type variables to their actual types. +newtype Subst = Subst {unSubst :: M.Map TypeName Type} + +subst :: Subst -> Type -> Type +subst sigma = \case + VarTy x -> + case M.lookup x (unSubst sigma) of + Nothing -> VarTy x + Just t -> subst sigma t + IntTy -> IntTy + BoolTy -> BoolTy + ListTy t -> ListTy (subst sigma t) + TupleTy ts -> TupleTy (map (subst sigma) ts) + CallableTy ts ret -> CallableTy (map (subst sigma) ts) (subst sigma ret) + StringTy -> StringTy + SideEffectTy -> SideEffectTy + +unifyTyVar :: (MonadState Subst m, MonadError Error m) => TypeName -> Type -> m () +unifyTyVar x t = + if x `elem` freeTyVars t + then throwTypeError $ "type equation loops: " ++ formatType (VarTy x) ++ " = " ++ formatType t + else do + modify' (Subst . M.insert x t . unSubst) -- This doesn't introduce the loop. + +unifyType :: (MonadState Subst m, MonadError Error m) => Type -> Type -> m () +unifyType t1 t2 = do + sigma <- get + t1 <- return $ subst sigma t1 -- shadowing + t2 <- return $ subst sigma t2 -- shadowing + case (t1, t2) of + _ | t1 == t2 -> return () + (VarTy x1, _) -> do + unifyTyVar x1 t2 + (_, VarTy x2) -> do + unifyTyVar x2 t1 + (ListTy t1, ListTy t2) -> do + unifyType t1 t2 + (TupleTy ts1, TupleTy ts2) -> do + if length ts1 == length ts2 + then mapM_ (uncurry unifyType) (zip ts1 ts2) + else throwTypeError $ "type " ++ formatType t1 ++ " is not type " ++ formatType t2 + (CallableTy args1 ret1, CallableTy args2 ret2) -> do + if length args1 == length args2 + then mapM_ (uncurry unifyType) (zip args1 args2) + else throwTypeError $ "type " ++ formatType t1 ++ " is not type " ++ formatType t2 + unifyType ret1 ret2 + _ -> throwTypeError $ "type " ++ formatType t1 ++ " is not type " ++ formatType t2 + +solveEquations :: MonadError Error m => [(Type, Type, Maybe Loc)] -> m Subst +solveEquations eqns = wrapError' "failed to solve type equations" $ do + flip execStateT (Subst M.empty) $ do + errs <- forM eqns $ \(t1, t2, loc) -> do + (Right <$> unifyType t1 t2) `catchError` \err -> do + sigma <- get + t1 <- return $ subst sigma t1 -- shadowing + t2 <- return $ subst sigma t2 -- shadowing + return $ Left (maybe id WithLocation loc (WithWrapped ("failed to unify type " ++ formatType t1 ++ " and type " ++ formatType t2) err)) + reportErrors errs + +mapTypeConstant :: (Type -> Type) -> Constant -> Constant +mapTypeConstant f = \case + ConstNone -> ConstNone + ConstInt n -> ConstInt n + ConstBool p -> ConstBool p + ConstBuiltin b -> ConstBuiltin (mapTypeBuiltin f b) + +mapTypeTarget :: (Type -> Type) -> Target' -> Target' +mapTypeTarget f = fmap $ \case + SubscriptTrg x index -> SubscriptTrg (mapTypeTarget f x) (mapTypeExpr f index) + NameTrg x -> NameTrg x + TupleTrg xs -> TupleTrg (map (mapTypeTarget f) xs) + +mapTypeExpr :: (Type -> Type) -> Expr' -> Expr' +mapTypeExpr f = mapSubExpr go + where + go = fmap $ \case + Lambda args body -> Lambda (map (second f) args) (go body) + ListComp e (Comprehension x iter pred) -> ListComp (go e) (Comprehension (mapTypeTarget f x) (go iter) (fmap go pred)) + Compare e1 (CmpOp' op t) e2 -> Compare (go e1) (CmpOp' op (f t)) (go e2) + Constant const -> Constant (mapTypeConstant f const) + Attribute e a -> Attribute (go e) (mapTypeAttribute f <$> a) + List t es -> List (f t) (map go es) + e -> e + +mapTypeStatement :: (Type -> Type) -> Statement -> Statement +mapTypeStatement f = \case + Return e -> Return (mapTypeExpr f e) + AugAssign x op e -> AugAssign (mapTypeTarget f x) op (mapTypeExpr f e) + AnnAssign x t e -> AnnAssign (mapTypeTarget f x) (f t) (mapTypeExpr f e) + For x iter body -> For (mapTypeTarget f x) (mapTypeExpr f iter) (map (mapTypeStatement f) body) + If pred body1 body2 -> If (mapTypeExpr f pred) (map (mapTypeStatement f) body1) (map (mapTypeStatement f) body2) + Assert e -> Assert (mapTypeExpr f e) + Expr' e -> Expr' (mapTypeExpr f e) + +mapTypeToplevelStatement :: (Type -> Type) -> ToplevelStatement -> ToplevelStatement +mapTypeToplevelStatement f = \case + ToplevelAnnAssign x t e -> ToplevelAnnAssign x (f t) (mapTypeExpr f e) + ToplevelFunctionDef g args ret body -> ToplevelFunctionDef g (map (second f) args) (f ret) (map (mapTypeStatement f) body) + ToplevelAssert e -> ToplevelAssert (mapTypeExpr f e) + +mapTypeProgram :: (Type -> Type) -> Program -> Program +mapTypeProgram f prog = map (mapTypeToplevelStatement f) prog + +-- | `substUnit` replaces all undetermined type variables with the unit type. +substUnit :: Type -> Type +substUnit = \case + VarTy _ -> NoneTy + IntTy -> IntTy + BoolTy -> BoolTy + ListTy t -> ListTy (substUnit t) + TupleTy ts -> TupleTy (map substUnit ts) + CallableTy ts ret -> CallableTy (map substUnit ts) (substUnit ret) + StringTy -> StringTy + SideEffectTy -> SideEffectTy + +-- | `subst'` does `subst` and replaces all undetermined type variables with the unit type. +subst' :: Subst -> Type -> Type +subst' sigma = substUnit . subst sigma + +-- | `run` infers types of given programs. +-- +-- As the interface, you can understand this function does the following: +-- +-- 1. Finds a type environment \(\Gamma\) s.t. for all statement \(\mathrm{stmt}\) in the given program, \(\Gamma \vdash \mathrm{stmt}\) holds, and +-- 2. Annotates each variable in the program using the \(\Gamma\). +-- +-- In its implementation, this is just something like a Hindley-Milner type inference. +-- +-- == Requirements +-- +-- * There must be no name conflicts in given programs. They must be alpha-converted. (`Jikka.RestrictedPython.Convert.Alpha`) +-- * All names must be resolved. (`Jikka.RestrictedPython.Convert.ResolveBuiltin`) +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.RestrictedPython.Convert.TypeInfer" $ do + eqns <- formularizeProgram prog + let (eqns', assertions) = sortEquations eqns + let eqns'' = mergeAssertions assertions + sigma <- solveEquations (eqns' ++ eqns'') + return $ mapTypeProgram (subst' sigma) prog diff --git a/src/Jikka/RestrictedPython/Convert/UseAppend.hs b/src/Jikka/RestrictedPython/Convert/UseAppend.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Convert/UseAppend.hs @@ -0,0 +1,51 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.RestrictedPython.Convert.UseAppend +-- Description : converts @xs = xs + [x]@ and @xs += [x]@ to @xs.append(x)@. / @xs = xs + [x]@ と @xs += [x]@ を @xs.append(x)@ に変換します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.RestrictedPython.Convert.UseAppend + ( run, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Util + +runStatement :: Statement -> [Statement] +runStatement = \case + AugAssign xs Add (WithLoc' _ (List t [x])) -> + [Expr' (withoutLoc (Call (withoutLoc (Attribute (targetToExpr xs) (withoutLoc (BuiltinAppend t)))) [x]))] + AnnAssign xs t1 (WithLoc' _ (BinOp xs' Add (WithLoc' _ (List t2 [x])))) + | dropLocation (targetToExpr xs) == dropLocation xs' -> + let t = case t1 of + ListTy t -> t + _ -> t2 + in [Expr' (withoutLoc (Call (withoutLoc (Attribute (targetToExpr xs) (withoutLoc (BuiltinAppend t)))) [x]))] + stmt -> [stmt] + +-- | `run` converts @xs = xs + [x]@ and @xs += [x]@ to @xs.append(x)@. +-- +-- == Examples +-- +-- Before: +-- +-- > xs = xs + [x] +-- > xs += [x] +-- > xs.append(x) +-- +-- After: +-- +-- > xs.append(x) +-- > xs.append(x) +-- > xs.append(x) +run :: (MonadAlpha m, MonadError Error m) => Program -> m Program +run prog = wrapError' "Jikka.RestrictedPython.Convert.UseAppend" $ do + return $ mapStatement runStatement prog diff --git a/src/Jikka/RestrictedPython/Evaluate.hs b/src/Jikka/RestrictedPython/Evaluate.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Evaluate.hs @@ -0,0 +1,682 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.RestrictedPython.Evaluate +-- Description : evaluates programs of the restricted Python. / 制限された Python のプログラムを評価します。 +-- Copyright : (c) Kimiyuki Onaka, 2021 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.RestrictedPython.Evaluate + ( run, + + -- * internal functions + makeGlobal, + runWithGlobal, + evalExpr, + evalStatement, + evalStatements, + execToplevelStatement, + ) +where + +import Control.Arrow (first) +import Control.Monad.Reader +import Control.Monad.State.Strict +import Data.Bits +import Data.List (maximumBy, minimumBy, sortBy) +import qualified Data.Map.Strict as M +import qualified Data.Vector as V +import Jikka.Common.Combinatorics +import Jikka.Common.Error +import Jikka.RestrictedPython.Format (formatAttribute, formatBuiltin, formatOperator) +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Lint +import Jikka.RestrictedPython.Language.Util +import Jikka.RestrictedPython.Language.Value + +assign :: MonadState Local m => VarName -> Value -> m () +assign x v = modify' (Local . M.insert x v . unLocal) + +lookupLocal :: (MonadState Local m, MonadError Error m) => VarName' -> m Value +lookupLocal x = do + local <- get + case M.lookup (value' x) (unLocal local) of + Just v -> return v + Nothing -> throwInternalErrorAt' (loc' x) $ "undefined variable: " ++ unVarName (value' x) + +assignSubscriptedTarget :: (MonadReader Global m, MonadState Local m, MonadError Error m) => Target' -> Expr' -> Value -> m () +assignSubscriptedTarget f index v = wrapAt' (loc' f) $ do + let go f indices = wrapAt' (loc' f) $ case value' f of + SubscriptTrg f index -> go f (index : indices) + NameTrg x -> return (x, indices) + TupleTrg _ -> throwInternalError "cannot subscript a tuple target" + (x, indices) <- go f [index] + f <- lookupLocal x + indices <- mapM evalExpr indices + let go f index = case (f, index) of + (_, []) -> return v + (ListVal f, IntVal index : indices) -> do + when (index < 0 || fromIntegral (V.length f) <= index) $ do + throwRuntimeError "list index out of range" + v' <- go (f V.! fromInteger index) indices + return $ ListVal (f V.// [(fromInteger index, v')]) + (_, _) -> throwInternalError "type error" + f <- go f indices + assign (value' x) f + +assignTarget :: (MonadReader Global m, MonadState Local m, MonadError Error m) => Target' -> Value -> m () +assignTarget x0 v = wrapAt' (loc' x0) $ case value' x0 of + SubscriptTrg f index -> do + assignSubscriptedTarget f index v + NameTrg x -> do + assign (value' x) v + TupleTrg xs -> do + case v of + TupleVal vs -> do + when (length xs /= length vs) $ do + throwInternalError "the lengths of tuple are different" + forM_ (zip xs vs) $ \(x, v) -> do + assignTarget x v + _ -> throwInternalError "cannot unpack non-tuple value" + +evalTarget :: (MonadReader Global m, MonadState Local m, MonadError Error m) => Target' -> m Value +evalTarget x0 = wrapAt' (loc' x0) $ case value' x0 of + SubscriptTrg f index -> do + f <- evalTarget f + index <- evalExpr index + case (f, index) of + (ListVal f, IntVal index) -> do + when (index < 0 || fromIntegral (V.length f) <= index) $ do + throwRuntimeError "list index out of range" + return $ f V.! fromInteger index + (_, _) -> throwInternalError "type error" + NameTrg x -> lookupLocal x + TupleTrg xs -> TupleVal <$> mapM evalTarget xs + +-- | `evalExpr` evaluates exprs of our restricted Python-like language. +-- +-- === Rules for \(e_1 \operatorname{boolop} e_2\) +-- +-- \[ +-- \cfrac{e_1 \mid \mu \Downarrow \mathbf{false}}{e_1 ~\mathbf{and}~ e_2 \mid \mu \Downarrow \mathbf{false}} +-- \] +-- \[ +-- \cfrac{e_1 \mid \mu \Downarrow \mathbf{true} \qquad e_2 \mid \mu \Downarrow v}{e_1 ~\mathbf{and}~ e_2 \mid \mu \Downarrow v} +-- \] +-- \[ +-- \vdots +-- \] +-- +-- === Rules for \(e_1 \operatorname{binop} e_2\) +-- +-- \[ +-- \cfrac{e_1 \mid \mu \Downarrow v_1 \qquad e_2 \mid \mu \Downarrow v_2}{e_1 + e_2 \mid \mu \Downarrow (v_1 + v_2)} +-- \] +-- \[ +-- \vdots +-- \] +-- +-- === Rules for \(\operatorname{unaryop} e\) +-- +-- === Rules for \(\lambda x _ \tau x _ \tau \dots x _ \tau. e\) +-- +-- \[ +-- \lambda {x_0} _ {\tau _ 0} {x_1} _ {\tau _ 1} \dots {x _ {n - 1}} _ {\tau _ {n - 1}}. e \mid \mu \Downarrow \lambda _ {\mu} x_0 x_1 \dots x _ {n - 1}. e +-- \] +-- +-- === Rules for \(\mathbf{if}~ e ~\mathbf{then}~ e ~\mathbf{else}~ e\) +-- +-- === Rules for \(\lbrack e ~\mathbf{for}~ y ~\mathbf{in}~ e ~(\mathbf{if}~ e)? \rbrack\) +-- +-- === Rules for \(e \operatorname{cmpop} e\) +-- +-- === Rules for \(e (e, e, \dots, e)\) +-- +-- \[ +-- \cfrac{ +-- e \mid \mu \Downarrow \lambda _ {\mu'} x_0 x_1 \dots x _ {n - 1}. e' +-- \qquad e_0 \mid \mu \Downarrow v_0 +-- \qquad e_1 \mid \mu \Downarrow v_1 +-- \qquad \dots +-- \qquad e _ {n - 1} \mid \mu \Downarrow v _ {n - 1} +-- \qquad e' \mid (x_0 \mapsto v_0; x_1 \mapsto v_1; \dots; x _ {n - 1} \mapsto v _ {n - 1}; \mu') \Downarrow v +-- }{ +-- e(e_0, e_1, \dots, e _ {n - 1}) \mid \mu \Downarrow v +-- } +-- \] +-- +-- \[ +-- \cfrac{ +-- e \mid \mu \Downarrow b +-- \qquad e_0 \mid \mu \Downarrow v_0 +-- \qquad e_1 \mid \mu \Downarrow v_1 +-- \qquad \dots +-- \qquad e _ {n - 1} \mid \mu \Downarrow v _ {n - 1} +-- }{ +-- e(e_0, e_1, \dots, e _ {n - 1}) \mid \mu \Downarrow b(e_0, e_1, \dots, e _ {n - 1}) +-- } +-- \qquad{(b ~\text{is a builtin function})} +-- \] +-- +-- === Rules for \(\operatorname{constant}\) +-- +-- === Rules for \(e \lbrack e \rbrack\) +-- +-- === Rules for \(x\) +-- +-- \[ +-- x \mid \mu \Downarrow \mu(x) +-- \] +-- +-- === Rules for \(\lbrack e, e, \dots, e \rbrack _ \tau\) +-- +-- === Rules for \(e \lbrack e? \colon e? \colon e? \rbrack\) +evalExpr :: (MonadReader Global m, MonadState Local m, MonadError Error m) => Expr' -> m Value +evalExpr e0 = wrapAt' (loc' e0) $ case value' e0 of + BoolOp e1 op e2 -> do + v1 <- evalExpr e1 + case (v1, op) of + (BoolVal False, And) -> return $ BoolVal False + (BoolVal True, And) -> evalExpr e2 + (BoolVal False, Or) -> evalExpr e2 + (BoolVal True, Or) -> return $ BoolVal True + (BoolVal False, Implies) -> return $ BoolVal True + (BoolVal True, Implies) -> evalExpr e2 + (_, _) -> throwInternalError "type error" + BinOp e1 op e2 -> do + v1 <- evalExpr e1 + v2 <- evalExpr e2 + evalBinOp v1 op v2 + UnaryOp op e -> do + v <- evalExpr e + case (op, v) of + (Invert, IntVal v) -> return $ IntVal (complement v) + (Not, BoolVal v) -> return $ BoolVal (not v) + (UAdd, IntVal v) -> return $ IntVal v + (USub, IntVal v) -> return $ IntVal (- v) + (_, _) -> throwInternalError "type error" + Lambda args body -> do + savedLocal <- get + return $ ClosureVal savedLocal (map (first value') args) [Return body] + IfExp e1 e2 e3 -> do + v1 <- evalExpr e1 + case v1 of + BoolVal True -> evalExpr e2 + BoolVal False -> evalExpr e3 + _ -> throwInternalError "type error" + ListComp e (Comprehension x iter pred) -> do + iter <- evalExpr iter + case iter of + ListVal iter -> do + savedLocal <- get + vs <- V.forM iter $ \it -> do + assignTarget x it + pred <- mapM evalExpr pred + case pred of + Just (BoolVal False) -> return Nothing + _ -> Just <$> evalExpr e + put savedLocal + return $ ListVal (V.catMaybes vs) + _ -> throwInternalError "type error" + Compare e1 op e2 -> do + v1 <- evalExpr e1 + v2 <- evalExpr e2 + case op of + CmpOp' In _ -> do + v2 <- toList v2 + return $ BoolVal (v1 `V.elem` v2) + CmpOp' NotIn _ -> do + v2 <- toList v2 + return $ BoolVal (v1 `V.elem` v2) + CmpOp' op _ -> do + ordering <- maybe (throwInternalError "something wrong") return (compareValues v1 v2) + BoolVal <$> case op of + Eq' -> return $ ordering == EQ + NotEq -> return $ ordering /= EQ + Lt -> return $ ordering == LT + LtE -> return $ ordering /= GT + Gt -> return $ ordering == GT + GtE -> return $ ordering /= LT + Is -> return $ ordering == EQ + IsNot -> return $ ordering /= EQ + _ -> throwInternalError "something wrong" + Call f args -> evalCall f args + Constant const -> + return $ case const of + ConstNone -> TupleVal [] + ConstInt v -> IntVal v + ConstBool v -> BoolVal v + ConstBuiltin v -> BuiltinVal v + Attribute e a -> AttributeVal <$> evalExpr e <*> pure (value' a) + Subscript e1 e2 -> do + v1 <- evalExpr e1 + v2 <- evalExpr e2 + case (v1, v2) of + (ListVal v1, IntVal v2) -> do + when (v2 < 0 || fromIntegral (V.length v1) <= v2) $ do + throwRuntimeError "list index out of range" + return $ v1 V.! fromInteger v2 + _ -> throwInternalError "type error" + Starred _ -> + throwInternalError "cannot evaluate starred expr" + Name x -> do + local <- get + case M.lookup (value' x) (unLocal local) of + Just v -> return v + Nothing -> do + global <- ask + case M.lookup (value' x) (unGlobal global) of + Just v -> return v + Nothing -> throwInternalError $ "undefined variable: " ++ unVarName (value' x) + List _ es -> ListVal . V.fromList <$> mapM evalExpr es + Tuple es -> TupleVal <$> mapM evalExpr es + SubscriptSlice e from to step -> do + v <- evalExpr e + from <- mapM evalExpr from + to <- mapM evalExpr to + step <- mapM evalExpr step + case v of + ListVal v -> + ListVal <$> case (from, to, step) of + (_, _, Just _) -> throwInternalError "slice with step is TODO" + (Nothing, Nothing, Nothing) -> return v + (Nothing, Just (IntVal to), Nothing) -> return $ V.take (fromInteger to) v + (Just (IntVal from), Nothing, Nothing) -> return $ V.drop (fromInteger from) v + (Just (IntVal from), Just (IntVal to), Nothing) -> return $ V.drop (fromInteger from) (V.take (fromInteger to) v) + (_, _, _) -> throwInternalError "type error" + _ -> throwInternalError "type error" + +evalCall :: (MonadReader Global m, MonadState Local m, MonadError Error m) => Expr' -> [Expr'] -> m Value +evalCall f args = wrapAt' (loc' f) $ do + f <- evalExpr f + args <- mapM evalExpr args + evalCall' f args + +evalCall' :: (MonadReader Global m, MonadState Local m, MonadError Error m) => Value -> [Value] -> m Value +evalCall' f actualArgs = case f of + AttributeVal v a -> do + evalAttribute v a actualArgs + BuiltinVal b -> do + evalBuiltin b actualArgs + ClosureVal local formalArgs body -> do + when (length formalArgs /= length actualArgs) $ do + throwInternalError "wrong number of arguments" + savedLocal <- get + put local + mapM_ (uncurry assign) (zip (map fst formalArgs) actualArgs) + v <- evalStatements body + put savedLocal + case v of + Just v -> return v + Nothing -> throwRuntimeError "it reaches the end of function without return" + _ -> throwRuntimeError "type error" + +-- | `evalStatement` evaluates statements of our restricted Python-like language. +-- When a statement is evaluated, it returns a value \(v\), doesn't return anything \(\mathbf{stop}\), or fails \(\mathbf{err}\). +-- Also it updates the environment function \(\mu\) from variables to values. +-- +-- === Rules for \(\mathbf{return}~ e\) +-- +-- \[ +-- \cfrac{ +-- e \mid \mu \Downarrow v +-- }{ +-- \mathbf{return}~ e \mid \mu \Downarrow v \mid \mu +-- } +-- \] +-- +-- === Rules for \(y \operatorname{binop} = e\) +-- +-- \[ +-- \cfrac{ +-- y \operatorname{binop} e \mid \mu \Downarrow v +-- }{ +-- y \operatorname{binop} = e \mid \mu \Downarrow \mathbf{stop} \mid (y \mapsto v; \mu) +-- } +-- \] +-- +-- === Rules for \(y := e\) +-- +-- \[ +-- \cfrac{ +-- e \mid \mu \Downarrow v +-- }{ +-- y \operatorname{binop} = e \mid \mu \Downarrow \mathbf{stop} \mid (y \mapsto v; \mu) +-- } +-- \] +-- +-- === Rules for \(\mathbf{for}~ y ~\mathbf{in}~ e \colon\quad \mathrm{stmt}; \mathrm{stmt}; \dots; \mathrm{stmt}\) +-- +-- \[ +-- \cfrac{ +-- e \mid \mu \Downarrow \mathbf{nil} +-- }{ +-- (\mathbf{for}~ y ~\mathbf{in}~ e \colon~ \vec{\mathrm{stmt}}) \mid \mu \Downarrow \mathbf{stop} \mid \mu +-- } +-- \] +-- +-- \[ +-- \cfrac{ +-- e \mid \mu \Downarrow \mathbf{cons}(v_1, v_2) +-- \qquad \vec{\mathrm{stmt}} \mid (y \mapsto v_1; \mu) \Downarrow v \mid \mu' +-- }{ +-- (\mathbf{for}~ y ~\mathbf{in}~ e \colon~ \vec{\mathrm{stmt}}) \mid \mu \Downarrow v \mid \mu' +-- } +-- \] +-- +-- \[ +-- \cfrac{ +-- e \mid \mu \Downarrow \mathbf{cons}(v_1, v_2) +-- \qquad \vec{\mathrm{stmt}} \mid (y \mapsto v_1; \mu) \Downarrow \mathbf{stop} \mid \mu' +-- \qquad (\mathbf{for}~ y ~\mathbf{in}~ v_2 \colon~ \vec{\mathrm{stmt}}) \mid \mu' \Downarrow a \mid \mu'' +-- }{ +-- (\mathbf{for}~ y ~\mathbf{in}~ e \colon~ \vec{\mathrm{stmt}}) \mid \mu \Downarrow a \mid \mu'' +-- } +-- \qquad{(a \in \lbrace v, \mathbf{stop} \rbrace)} +-- \] +-- +-- It assumes the program is properly alpha-converted, i.e. `doesntHaveLeakOfLoopCounters`. So it leaks loop counters to out of loops. +-- +-- === Rules for \(\mathbf{if}~ e \colon\quad \mathrm{stmt}; \mathrm{stmt}; \dots; \mathrm{stmt};\quad \mathbf{else}\colon\quad \mathrm{stmt}; \mathrm{stmt}; \dots; \mathrm{stmt}\) +-- +-- \[ +-- \cfrac{ +-- e \mid \mu \Downarrow \mathbf{true} +-- \qquad \vec{\mathrm{stmt}} _ 1 \mid \mu \Downarrow a \mid \mu' +-- }{ +-- (\mathbf{if}~ e \colon~ \vec{\mathrm{stmt}} _ 1 ~\mathbf{else}\colon~ \vec{\mathrm{stmt}} _ 2) \mid \mu \Downarrow a \mid \mu' +-- } +-- \qquad{(a \in \lbrace v, \mathbf{stop} \rbrace)} +-- \] +-- +-- \[ +-- \cfrac{ +-- e \mid \mu \Downarrow \mathbf{false} +-- \qquad \vec{\mathrm{stmt}} _ 2 \mid \mu \Downarrow a \mid \mu' +-- }{ +-- (\mathbf{if}~ e \colon~ \vec{\mathrm{stmt}} _ 1 ~\mathbf{else}\colon~ \vec{\mathrm{stmt}} _ 2) \mid \mu \Downarrow a \mid \mu' +-- } +-- \qquad{(a \in \lbrace v, \mathbf{stop} \rbrace)} +-- \] +-- +-- === Rules for \(\mathbf{assert}~ e\) +-- +-- \[ +-- \cfrac{ +-- e \mid \mu \Downarrow \mathbf{true} +-- }{ +-- \mathbf{assert}~ e \mid \mu \Downarrow \mathbf{stop} +-- } +-- \] +-- +-- \[ +-- \cfrac{ +-- e \mid \mu \Downarrow \mathbf{false} +-- }{ +-- \mathbf{assert}~ e \mid \mu \Downarrow \mathbf{err} +-- } +-- \] +-- +-- === Rules for \(e\) +-- +-- \[ +-- \cfrac{ +-- e \mid \mu \Downarrow v +-- }{ +-- x.\mathrm{append}(e) \mid \mu \Downarrow \mathbf{stop} \mid (x \mapsto \mathrm{snoc}(\mu(x), v); \mu) +-- } +-- \] +evalStatement :: (MonadReader Global m, MonadState Local m, MonadError Error m) => Statement -> m (Maybe Value) +evalStatement = \case + Return e -> do + v <- evalExpr e + return $ Just v + AugAssign x op e -> do + v1 <- evalTarget x + v2 <- evalExpr e + v <- evalBinOp v1 op v2 + assignTarget x v + return Nothing + AnnAssign x _ e -> do + v <- evalExpr e + assignTarget x v + return Nothing + For x iter body -> do + iter <- evalExpr iter + case iter of + ListVal iter -> do + let go [] = return Nothing + go (it : iter) = do + assignTarget x it + v <- evalStatements body + case v of + Just v -> return $ Just v + Nothing -> go iter + go (V.toList iter) + _ -> wrapAt' (loc' x) $ do + throwInternalError "type error" + If pred body1 body2 -> do + pred <- evalExpr pred + if pred /= BoolVal False + then evalStatements body1 + else evalStatements body2 + Assert e -> wrapAt' (loc' e) $ do + v <- evalExpr e + when (v == BoolVal False) $ do + throwRuntimeError "assertion failure" + return Nothing + Append loc _ x e -> case exprToTarget x of + Nothing -> throwSemanticErrorAt' loc "wrong `append` method" + Just x -> do + v1 <- evalTarget x + v2 <- evalExpr e + v <- ListVal <$> (V.snoc <$> toList v1 <*> pure v2) + assignTarget x v + return Nothing + Expr' e -> throwSemanticErrorAt' (loc' e) "wrong expr-statement" + +-- | `evalStatements` evaluates sequences of statements of our restricted Python-like language. +-- +-- \[ +-- \cfrac{\mathrm{stmt} _ 0 \mid \mu \Downarrow v \mid \mu'}{\mathrm{stmt} _ 0; \mathrm{stmt} _ 1; \dots; \mathrm{stmt} _ {n-1} \mid \mu \Downarrow v \mid \mu'} +-- \] +-- +-- \[ +-- \cfrac{\mathrm{stmt} _ 0 \mid \mu \Downarrow \mathbf{stop} \mid \mu' \qquad \mathrm{stmt} _ 1; \dots; \mathrm{stmt} _ {n-1} \mid \mu' \Downarrow a \mid \mu''}{\mathrm{stmt} _ 0; \mathrm{stmt} _ 1; \dots; \mathrm{stmt} _ {n-1} \mid \mu \Downarrow a \mid \mu''} +-- \qquad{(a \in \lbrace v, \mathbf{stop} \rbrace)} +-- \] +-- +-- \[ +-- \epsilon \mid \mu \Downarrow \mathbf{stop} \mid \mu +-- \] +evalStatements :: (MonadReader Global m, MonadState Local m, MonadError Error m) => [Statement] -> m (Maybe Value) +evalStatements [] = return Nothing +evalStatements (stmt : stmts) = do + v <- evalStatement stmt + case v of + Just v -> return $ Just v + Nothing -> evalStatements stmts + +execToplevelStatement :: (MonadState Global m, MonadError Error m) => ToplevelStatement -> m () +execToplevelStatement = \case + ToplevelAnnAssign x _ e -> do + global <- get + v <- runWithGlobal global e + put $ Global (M.insert (value' x) v (unGlobal global)) + ToplevelFunctionDef f args _ body -> do + global <- get + let v = ClosureVal (Local M.empty) (map (first value') args) body + put $ Global (M.insert (value' f) v (unGlobal global)) + ToplevelAssert e -> do + global <- get + v <- runWithGlobal global e + when (v /= BoolVal True) $ do + throwRuntimeError "assertion failure" + +newtype Global = Global + { unGlobal :: M.Map VarName Value + } + deriving (Eq, Ord, Show, Read) + +initialGlobal :: Global +initialGlobal = Global M.empty + +lookupGlobal :: MonadError Error m => VarName' -> Global -> m Value +lookupGlobal x global = + case M.lookup (value' x) (unGlobal global) of + Just y -> return y + Nothing -> throwSymbolErrorAt' (loc' x) $ "undefined variable: " ++ unVarName (value' x) + +runWithGlobal :: MonadError Error m => Global -> Expr' -> m Value +runWithGlobal global e = do + runReaderT (evalStateT (evalExpr e) (Local M.empty)) global + +runWithGlobal' :: MonadError Error m => Global -> Value -> [Value] -> m Value +runWithGlobal' global solve args = do + runReaderT (evalStateT (evalCall' solve args) (Local M.empty)) global + +-- | `makeGlobal` packs toplevel definitions into `Global`. +-- This assumes `doesntHaveLeakOfLoopCounters`. +makeGlobal :: MonadError Error m => Program -> m Global +makeGlobal prog = do + ensureDoesntHaveLeakOfLoopCounters prog + execStateT (mapM_ execToplevelStatement prog) initialGlobal + +run :: MonadError Error m => Program -> [Value] -> m Value +run prog args = wrapError' "Jikka.RestrictedPython.Evaluate" $ do + global <- makeGlobal prog + solve <- lookupGlobal (withoutLoc (VarName "solve")) global + runWithGlobal' global solve args + +evalBinOp :: MonadError Error m => Value -> Operator -> Value -> m Value +evalBinOp v1 op v2 = wrapError' ("calculating " ++ formatOperator op ++ " operator") $ do + v1 <- toInt v1 + v2 <- toInt v2 + v <- case (op, v2) of + (Add, _) -> return $ v1 + v2 + (Sub, _) -> return $ v1 - v2 + (Mult, _) -> return $ v1 * v2 + (MatMult, _) -> throwInternalError "matmul operator ('@') is not supported" + (Div, _) -> throwInternalError "floatdiv operator ('/') is not supported" + (FloorDiv, 0) -> throwRuntimeError "division by zero" + (FloorDiv, _) -> return $ v1 `div` v2 + (FloorMod, 0) -> throwRuntimeError "division by zero" + (FloorMod, _) -> return $ v1 `mod` v2 + (CeilDiv, 0) -> throwRuntimeError "division by zero" + (CeilDiv, _) -> return $ (v1 + v2 - 1) `div` v2 + (CeilMod, 0) -> throwRuntimeError "division by zero" + (CeilMod, _) -> return $ (v1 + v2 - 1) `mod` v2 + (Pow, _) -> return $ v1 ^ v2 + (BitLShift, _) -> return $ shiftL v1 (fromInteger v2) + (BitRShift, _) -> return $ shiftR v1 (fromInteger v2) + (BitOr, _) -> return $ v1 .|. v2 + (BitXor, _) -> return $ v1 `xor` v2 + (BitAnd, _) -> return $ v1 .&. v2 + (Max, _) -> return $ max v1 v2 + (Min, _) -> return $ min v1 v2 + return $ IntVal v + +evalBuiltin :: (MonadReader Global m, MonadState Local m, MonadError Error m) => Builtin -> [Value] -> m Value +evalBuiltin b args = wrapError' ("calling " ++ formatBuiltin b) $ do + let go1' t1 ret f = case args of + [v1] -> ret <$> (f =<< t1 v1) + _ -> throwInternalError $ "expected 1 argument, got " ++ show (length args) + let go1 t1 ret f = go1' t1 ret (return . f) + let go2' t1 t2 ret f = case args of + [v1, v2] -> ret <$> join (f <$> t1 v1 <*> t2 v2) + _ -> throwInternalError $ "expected 2 arguments, got " ++ show (length args) + let go2 t1 t2 ret f = go2' t1 t2 ret ((return .) . f) + let go3 t1 t2 t3 ret f = case args of + [v1, v2, v3] -> ret <$> (f <$> t1 v1 <*> t2 v2 <*> t3 v3) + _ -> throwInternalError $ "expected 3 arguments, got " ++ show (length args) + let goN' t ret f = ret <$> (f =<< mapM t args) + let goN t ret f = goN' t ret (return . f) + let zipN acc [] = reverse acc + zipN acc xss | any null xss = reverse acc + zipN acc xss = zipN (map head xss : acc) (map tail xss) + case b of + BuiltinAbs -> go1 toInt IntVal abs + BuiltinPow -> go2 toInt toInt IntVal (^) + BuiltinModPow -> go3 toInt toInt toInt IntVal $ \x k m -> (x ^ k) `mod` m + BuiltinAll -> go1 toBoolList BoolVal minimum + BuiltinAny -> go1 toBoolList BoolVal maximum + BuiltinDivMod -> go2' toInt toInt TupleVal $ \a b -> case b of + 0 -> throwRuntimeError "division by zero" + _ -> return [IntVal (a `div` b), IntVal (a `mod` b)] + BuiltinSorted _ -> go1 toList ListVal (V.fromList . sortBy compareValues' . V.toList) + BuiltinEnumerate _ -> go1 toList ListVal (V.fromList . zipWith (\i x -> TupleVal [IntVal i, x]) [0 ..] . V.toList) + BuiltinBool _ -> go1' pure BoolVal $ \case + IntVal n -> return $ n /= 0 + BoolVal p -> return p + ListVal xs -> return $ not (V.null xs) + TupleVal xs -> return $ not (null xs) + _ -> throwInternalError "type error" + BuiltinInt _ -> go1' return IntVal $ \case + IntVal n -> return n + BoolVal p -> return $ if p then 1 else 0 + _ -> throwInternalError "type error" + BuiltinTuple _ -> goN pure TupleVal id + BuiltinSum -> go1 toIntList IntVal sum + BuiltinZip _ -> goN toList ListVal (V.fromList . map TupleVal . zipN [] . map V.toList) + BuiltinFilter _ -> go2' pure toList ListVal $ \f xs -> do + let go x = do + pred <- evalCall' f [x] + case pred of + BoolVal True -> return $ Just x + BoolVal False -> return Nothing + _ -> throwInternalError "type error" + V.mapMaybeM go xs + BuiltinLen _ -> go1 toList IntVal (fromIntegral . V.length) + BuiltinList _ -> go1 toList ListVal id + BuiltinRange1 -> go1 toInt ListVal $ \to -> V.fromList (map IntVal [0 .. to - 1]) + BuiltinRange2 -> go2 toInt toInt ListVal $ \from to -> V.fromList (map IntVal [from .. to - 1]) + BuiltinRange3 -> go3 toInt toInt toInt ListVal $ \from to step -> V.fromList (map IntVal [from, from + step .. to - 1]) + BuiltinMap _ _ -> goN' pure ListVal $ \case + [] -> throwInternalError "type error" + f : args -> do + args <- mapM toList args + V.fromList <$> mapM (evalCall' f) (zipN [] (map V.toList args)) + BuiltinReversed _ -> go1 toList ListVal V.reverse + BuiltinMin1 _ -> go1 toList id (minimumBy compareValues') + BuiltinMin _ _ -> goN pure id (minimumBy compareValues') + BuiltinMax1 _ -> go1 toList id (maximumBy compareValues') + BuiltinMax _ _ -> goN pure id (maximumBy compareValues') + BuiltinArgMax _ -> go1 toList IntVal $ \xs -> snd (maximumBy (\(x, i) (y, j) -> compareValues' x y <> compare i j) (zip (V.toList xs) [0 ..])) + BuiltinArgMin _ -> go1 toList IntVal $ \xs -> snd (minimumBy (\(x, i) (y, j) -> compareValues' x y <> compare i j) (zip (V.toList xs) [0 ..])) + BuiltinCeilDiv -> go2' toInt toInt IntVal $ \a b -> if b == 0 then throwRuntimeError "division by zero" else return $ (a + b - 1) `div` b + BuiltinCeilMod -> go2' toInt toInt IntVal $ \a b -> if b == 0 then throwRuntimeError "division by zero" else return $ (a + b - 1) `mod` b + BuiltinFloorDiv -> go2' toInt toInt IntVal $ \a b -> if b == 0 then throwRuntimeError "division by zero" else return $ a `div` b + BuiltinFloorMod -> go2' toInt toInt IntVal $ \a b -> if b == 0 then throwRuntimeError "division by zero" else return $ a `mod` b + BuiltinGcd -> go2 toInt toInt IntVal gcd + BuiltinLcm -> go2 toInt toInt IntVal lcm + BuiltinModInv -> go2' toInt toInt IntVal $ \_ _ -> throwInternalError "Jikka.RestrictedPython.Evaluate.evalBuiltin: TODO" + BuiltinProduct -> go1 toIntList IntVal product + BuiltinFact -> go1' toInt IntVal $ \n -> if 0 <= n then return $ fact n else throwRuntimeError "invalid argument" + BuiltinChoose -> go2' toInt toInt IntVal $ \n r -> if 0 <= r && r <= n then return $ choose n r else throwRuntimeError "invalid argument" + BuiltinPermute -> go2' toInt toInt IntVal $ \n r -> if 0 <= r && r <= n then return $ permute n r else throwRuntimeError "invalid argument" + BuiltinMultiChoose -> go2' toInt toInt IntVal $ \n r -> if 0 <= r && r <= n then return $ multichoose n r else throwRuntimeError "invalid argument" + BuiltinInput -> throwSemanticError "cannot use `input' out of main function" + BuiltinPrint _ -> throwSemanticError "cannot use `print' out of main function" + +evalAttribute :: (MonadReader Global m, MonadState Local m, MonadError Error m) => Value -> Attribute -> [Value] -> m Value +evalAttribute v0 a args = wrapError' ("calling " ++ formatAttribute a) $ do + let go0' t0 ret f = case args of + [] -> ret <$> (f =<< t0 v0) + _ -> throwInternalError $ "expected 0 arguments, got " ++ show (length args) + let go0 t0 ret f = go0' t0 ret (return . f) + let go1' t0 t1 ret f = case args of + [v1] -> ret <$> join (f <$> t0 v0 <*> t1 v1) + _ -> throwInternalError $ "expected 1 argument, got " ++ show (length args) + let go1 t0 t1 ret f = go1' t0 t1 ret ((return .) . f) + case a of + UnresolvedAttribute a -> throwInternalError $ "Jikka.RestrictedPython.Evaluate.evalAttribute: unresolved attribute: " ++ unAttributeName a + BuiltinCount _ -> go1 toList pure IntVal $ \xs x -> toInteger (V.length (V.filter (== x) xs)) + BuiltinIndex _ -> go1' toList pure IntVal $ \xs x -> case V.elemIndex x xs of + Nothing -> throwRuntimeError $ "not in list: " ++ formatValue x + Just i -> return (toInteger i) + BuiltinCopy _ -> go0 toList ListVal id + BuiltinAppend _ -> throwSemanticError "cannot use `append' out of expr-statements" + BuiltinSplit -> throwSemanticError "cannot use `split' out of main function" diff --git a/src/Jikka/RestrictedPython/Format.hs b/src/Jikka/RestrictedPython/Format.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Format.hs @@ -0,0 +1,171 @@ +{-# LANGUAGE LambdaCase #-} + +-- | +-- Module : Jikka.RestrictedPython.Format +-- Description : converts AST of the restricted Python to strings. / 制限された Python の抽象構文木を文字列に変換します。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- TODO: add parens with considering precedences. +module Jikka.RestrictedPython.Format + ( run, + run', + formatType, + formatOperator, + formatBuiltin, + formatAttribute, + formatExpr, + formatTarget, + ) +where + +import Data.List (intercalate) +import Data.Text (Text, pack) +import Jikka.Common.Format.AutoIndent +import Jikka.RestrictedPython.Language.Builtin +import Jikka.RestrictedPython.Language.Expr + +formatType :: Type -> String +formatType t = case t of + VarTy x -> unTypeName x + IntTy -> "int" + BoolTy -> "bool" + ListTy t -> "List[" ++ formatType t ++ "]" + NoneTy -> "None" + TupleTy ts -> "Tuple[" ++ intercalate ", " (map formatType ts) ++ "]" + CallableTy ts ret -> "Callable[[" ++ intercalate ", " (map formatType ts) ++ "], " ++ formatType ret ++ "]" + StringTy -> "str" + SideEffectTy -> "side-effect" + +formatConstant :: Constant -> String +formatConstant = \case + ConstNone -> "None" + ConstInt n -> show n + ConstBool p -> show p + ConstBuiltin b -> formatBuiltin b + +formatBoolOp :: BoolOp -> String +formatBoolOp = \case + And -> "and" + Or -> "or" + Implies -> "implies" + +formatOperator :: Operator -> String +formatOperator = \case + Add -> "+" + Sub -> "-" + Mult -> "*" + MatMult -> "@" + Div -> "/" + FloorDiv -> "//" + FloorMod -> "%" + CeilDiv -> "/^" + CeilMod -> "%^" + Pow -> "**" + BitLShift -> "<<" + BitRShift -> ">>" + BitOr -> "|" + BitXor -> "^" + BitAnd -> "&" + Max -> ">?" + Min -> " String +formatUnaryOp = \case + Invert -> "~" + Not -> "not" + UAdd -> "+" + USub -> "-" + +formatCmpOp :: CmpOp' -> String +formatCmpOp (CmpOp' op _) = case op of + Eq' -> "==" + NotEq -> "!=" + Lt -> "<" + LtE -> "<=" + Gt -> ">" + GtE -> ">=" + Is -> "is" + IsNot -> "is not" + In -> "in" + NotIn -> "not in" + +formatComprehension :: Comprehension -> String +formatComprehension (Comprehension x iter ifs) = + let body = "for " ++ formatTarget x ++ " in " ++ formatExpr iter + ifs' = case ifs of + Nothing -> "" + Just ifs -> " if " ++ formatExpr ifs + in body ++ ifs' + +formatTarget :: Target' -> String +formatTarget (WithLoc' _ x) = case x of + SubscriptTrg x e -> formatTarget x ++ "[" ++ formatExpr e ++ "]" + NameTrg x -> unVarName (value' x) + TupleTrg xs -> case xs of + [] -> "()" + [x] -> "(" ++ formatTarget x ++ ",)" + _ -> intercalate ", " (map formatTarget xs) + +formatExpr :: Expr' -> String +formatExpr (WithLoc' _ e0) = case e0 of + BoolOp e1 op e2 -> formatExpr e1 ++ " " ++ formatBoolOp op ++ " " ++ formatExpr e2 + BinOp e1 op e2 -> formatExpr e1 ++ " " ++ formatOperator op ++ " " ++ formatExpr e2 + UnaryOp op e -> formatUnaryOp op ++ " " ++ formatExpr e + Lambda args body -> case args of + [] -> "lambda: " ++ formatExpr body + _ -> "lambda " ++ intercalate ", " (map (unVarName . value' . fst) args) ++ ": " ++ formatExpr body + IfExp e1 e2 e3 -> formatExpr e2 ++ " if " ++ formatExpr e1 ++ " else " ++ formatExpr e3 + ListComp e comp -> "[" ++ formatExpr e ++ " " ++ formatComprehension comp ++ "]" + Compare e1 op e2 -> formatExpr e1 ++ " " ++ formatCmpOp op ++ " " ++ formatExpr e2 + Call f args -> case args of + [WithLoc' _ (ListComp e comp)] -> formatExpr f ++ "(" ++ formatExpr e ++ " " ++ formatComprehension comp ++ ")" + _ -> formatExpr f ++ "(" ++ intercalate ", " (map formatExpr args) ++ ")" + Constant const -> formatConstant const + Attribute e (WithLoc' _ x) -> formatExpr e ++ "." ++ formatAttribute x + Subscript e1 e2 -> formatExpr e1 ++ "[" ++ formatExpr e2 ++ "]" + Starred e -> "*" ++ formatExpr e + Name x -> unVarName (value' x) + List _ es -> "[" ++ intercalate ", " (map formatExpr es) ++ "]" + Tuple es -> case es of + [] -> "()" + [e] -> "(" ++ formatExpr e ++ ",)" + _ -> "(" ++ intercalate ", " (map formatExpr es) ++ ")" + SubscriptSlice e from to step -> + let from' = maybe "" formatExpr from + to' = maybe "" formatExpr to + step' = maybe "" ((':' :) . formatExpr) step + in formatExpr e ++ "[" ++ from' ++ ":" ++ to' ++ step' ++ "]" + +formatStatement :: Statement -> [String] +formatStatement = \case + Return e -> ["return " ++ formatExpr e] + AugAssign x op e -> [formatTarget x ++ " " ++ formatOperator op ++ "= " ++ formatExpr e] + AnnAssign x t e -> [formatTarget x ++ ": " ++ formatType t ++ " = " ++ formatExpr e] + For x iter body -> ["for " ++ formatTarget x ++ " in " ++ formatExpr iter ++ ":", indent] ++ concatMap formatStatement body ++ [dedent] + If e body1 body2 -> case body2 of + [] -> ["if " ++ formatExpr e ++ ":", indent] ++ concatMap formatStatement body1 ++ [dedent] + [body2@(If _ _ _)] -> + let elif : cont = formatStatement body2 + in ["if " ++ formatExpr e ++ ":", indent] ++ concatMap formatStatement body1 ++ [dedent, "el" ++ elif] ++ cont + _ -> ["if " ++ formatExpr e ++ ":", indent] ++ concatMap formatStatement body1 ++ [dedent, "else:", indent] ++ concatMap formatStatement body2 ++ [dedent] + Assert e -> ["assert " ++ formatExpr e] + Expr' e -> [formatExpr e] + +formatToplevelStatement :: ToplevelStatement -> [String] +formatToplevelStatement = \case + ToplevelAnnAssign x t e -> [unVarName (value' x) ++ ": " ++ formatType t ++ " = " ++ formatExpr e] + ToplevelFunctionDef f args ret body -> ["def " ++ unVarName (value' f) ++ "(" ++ intercalate ", " (map (\(x, t) -> unVarName (value' x) ++ ": " ++ formatType t) args) ++ ") -> " ++ formatType ret ++ ":", indent] ++ concatMap formatStatement body ++ [dedent] + ToplevelAssert e -> ["assert " ++ formatExpr e] + +formatProgram :: Program -> [String] +formatProgram prog = concatMap formatToplevelStatement prog + +run' :: Program -> String +run' = unlines . makeIndentFromMarkers 4 . formatProgram + +run :: Applicative m => Program -> m Text +run = pure . pack . run' diff --git a/src/Jikka/RestrictedPython/Language/Builtin.hs b/src/Jikka/RestrictedPython/Language/Builtin.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Language/Builtin.hs @@ -0,0 +1,329 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.RestrictedPython.Language.Builtin where + +import Data.Functor +import qualified Data.Set as S +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Util + +builtinNames :: S.Set VarName +builtinNames = S.union standardBuiltinNames additionalBuiltinNames + +standardBuiltinNames :: S.Set VarName +standardBuiltinNames = + S.fromList + [ "abs", + "all", + "any", + "bool", + "divmod", + "enumerate", + "filter", + "input", + "int", + "len", + "list", + "map", + "max", + "min", + "pow", + "print", + "range", + "reversed", + "sorted", + "sum", + "zip" + ] + +additionalBuiltinNames :: S.Set VarName +additionalBuiltinNames = + S.fromList + [ "argmax", + "argmin", + "ceildiv", + "ceilmod", + "choose", + "fact", + "floordiv", + "floormod", + "gcd", + "inv", + "lcm", + "multichoose", + "permute", + "product" + ] + +-- | `resolveUniqueBuiltin` makes a builtin function from a variable name. +-- However, this doesn't anything for ambiguous builtin functions. +-- For example, the builtin function "max" is kept as a variable because it may be \(\mathbf{list}(\alpha) \to \alpha\), \(\alpha \times \alpha \to \alpha\), etc. and this function cannot resolve it. +resolveUniqueBuiltin :: (MonadAlpha m, MonadError Error m) => VarName' -> m Expr' +resolveUniqueBuiltin x | value' x `S.notMember` builtinNames = return $ WithLoc' (loc' x) (Name x) +resolveUniqueBuiltin x = do + let f = return . WithLoc' (loc' x) . Constant . ConstBuiltin + case value' x of + "abs" -> f BuiltinAbs + "all" -> f BuiltinAll + "any" -> f BuiltinAny + "bool" -> f . BuiltinBool =<< genType + "divmod" -> f BuiltinDivMod + "enumerate" -> f . BuiltinEnumerate =<< genType + "filter" -> f . BuiltinFilter =<< genType + "int" -> f . BuiltinInt =<< genType + "input" -> f BuiltinInput + "len" -> f . BuiltinLen =<< genType + "list" -> f . BuiltinList =<< genType + "reversed" -> f . BuiltinReversed =<< genType + "sorted" -> f . BuiltinSorted =<< genType + "sum" -> f BuiltinSum + "argmax" -> f . BuiltinArgMax =<< genType + "argmin" -> f . BuiltinArgMin =<< genType + "ceildiv" -> f BuiltinCeilDiv + "ceilmod" -> f BuiltinCeilMod + "choose" -> f BuiltinChoose + "fact" -> f BuiltinFact + "floordiv" -> f BuiltinFloorDiv + "floormod" -> f BuiltinFloorMod + "gcd" -> f BuiltinGcd + "inv" -> f BuiltinModInv + "lcm" -> f BuiltinLcm + "multichoose" -> f BuiltinMultiChoose + "permute" -> f BuiltinPermute + "product" -> f BuiltinProduct + _ -> return $ WithLoc' (loc' x) (Name x) + +resolveBuiltin :: (MonadAlpha m, MonadError Error m) => VarName' -> Int -> m Expr' +resolveBuiltin x _ | value' x `S.notMember` builtinNames = return $ WithLoc' (loc' x) (Name x) +resolveBuiltin x n = wrapAt' (loc' x) . wrapError' "Jikka.RestrictedPython.Language.Builtin.resolveBuiltin" $ do + let f = return . WithLoc' (loc' x) . Constant . ConstBuiltin + when (n < 0) $ do + throwInternalError $ "negative arity: " ++ show n + case value' x of + "map" -> f =<< (BuiltinMap <$> replicateM (n - 1) genType <*> genType) + "max" -> case n of + 1 -> f . BuiltinMax1 =<< genType + _ -> f =<< (BuiltinMax <$> genType <*> pure n) + "min" -> case n of + 1 -> f . BuiltinMin1 =<< genType + _ -> f =<< (BuiltinMin <$> genType <*> pure n) + "pow" -> + if n == 3 + then f BuiltinModPow + else f BuiltinPow + "print" -> f . BuiltinPrint =<< replicateM n genType + "range" -> case n of + 1 -> f BuiltinRange1 + 2 -> f BuiltinRange2 + 3 -> f BuiltinRange3 + _ -> throwTypeError $ "range expected 1, 2, or 3 arguments, got " ++ show n + "zip" -> f . BuiltinZip =<< replicateM n genType + _ -> do + e <- resolveUniqueBuiltin x + case value' e of + Constant (ConstBuiltin _) -> return e + _ -> throwInternalError $ "not exhaustive: " ++ unVarName (value' x) + +formatBuiltin :: Builtin -> String +formatBuiltin = \case + BuiltinAbs -> "abs" + BuiltinPow -> "pow" + BuiltinModPow -> "pow" + BuiltinAll -> "all" + BuiltinAny -> "any" + BuiltinDivMod -> "divmod" + BuiltinSorted _ -> "sorted" + BuiltinEnumerate _ -> "enumerate" + BuiltinBool _ -> "bool" + BuiltinInt _ -> "int" + BuiltinSum -> "sum" + BuiltinZip _ -> "zip" + BuiltinFilter _ -> "filter" + BuiltinTuple _ -> "tuple" + BuiltinLen _ -> "len" + BuiltinList _ -> "list" + BuiltinRange1 -> "range" + BuiltinRange2 -> "range" + BuiltinRange3 -> "range" + BuiltinMap _ _ -> "map" + BuiltinReversed _ -> "reversed" + BuiltinMax1 _ -> "max" + BuiltinMax _ _ -> "max" + BuiltinMin1 _ -> "min" + BuiltinMin _ _ -> "min" + BuiltinArgMax _ -> "argmax" + BuiltinArgMin _ -> "argmin" + BuiltinCeilDiv -> "ceildiv" + BuiltinCeilMod -> "ceilmod" + BuiltinFloorDiv -> "floordiv" + BuiltinFloorMod -> "floormod" + BuiltinChoose -> "choose" + BuiltinFact -> "fact" + BuiltinGcd -> "gcd" + BuiltinLcm -> "lcm" + BuiltinModInv -> "inv" + BuiltinMultiChoose -> "multichoose" + BuiltinPermute -> "permute" + BuiltinProduct -> "product" + BuiltinInput -> "input" + BuiltinPrint _ -> "print" + +typeBuiltin :: Builtin -> Type +typeBuiltin = \case + BuiltinAbs -> CallableTy [IntTy] IntTy + BuiltinPow -> CallableTy [IntTy] IntTy + BuiltinModPow -> CallableTy [IntTy, IntTy] IntTy + BuiltinAll -> CallableTy [ListTy BoolTy] BoolTy + BuiltinAny -> CallableTy [ListTy BoolTy] BoolTy + BuiltinArgMax t -> CallableTy [ListTy t] IntTy + BuiltinArgMin t -> CallableTy [ListTy t] IntTy + BuiltinBool t -> CallableTy [t] BoolTy + BuiltinCeilDiv -> CallableTy [IntTy, IntTy] IntTy + BuiltinCeilMod -> CallableTy [IntTy, IntTy] IntTy + BuiltinChoose -> CallableTy [IntTy, IntTy] IntTy + BuiltinDivMod -> CallableTy [IntTy, IntTy] (TupleTy [IntTy, IntTy]) + BuiltinEnumerate t -> CallableTy [ListTy t] (ListTy (TupleTy [IntTy, t])) + BuiltinFact -> CallableTy [ListTy IntTy] IntTy + BuiltinFilter t -> CallableTy [CallableTy [t] BoolTy, ListTy t] (ListTy t) + BuiltinFloorDiv -> CallableTy [IntTy, IntTy] IntTy + BuiltinFloorMod -> CallableTy [IntTy, IntTy] IntTy + BuiltinGcd -> CallableTy [IntTy, IntTy] IntTy + BuiltinInt t -> CallableTy [t] IntTy + BuiltinModInv -> CallableTy [IntTy, IntTy] IntTy + BuiltinLcm -> CallableTy [IntTy, IntTy] IntTy + BuiltinLen t -> CallableTy [ListTy t] IntTy + BuiltinList t -> CallableTy [ListTy t] (ListTy t) + BuiltinMap args ret -> CallableTy (CallableTy args ret : map ListTy args) (ListTy ret) + BuiltinMax t n -> CallableTy (replicate n t) t + BuiltinMax1 t -> CallableTy [ListTy t] t + BuiltinMin t n -> CallableTy (replicate n t) t + BuiltinMin1 t -> CallableTy [ListTy t] t + BuiltinMultiChoose -> CallableTy [IntTy, IntTy] IntTy + BuiltinPermute -> CallableTy [IntTy, IntTy] IntTy + BuiltinProduct -> CallableTy [ListTy IntTy] IntTy + BuiltinRange1 -> CallableTy [IntTy] (ListTy IntTy) + BuiltinRange2 -> CallableTy [IntTy, IntTy] (ListTy IntTy) + BuiltinRange3 -> CallableTy [IntTy, IntTy, IntTy] (ListTy IntTy) + BuiltinReversed t -> CallableTy [ListTy t] (ListTy t) + BuiltinSorted t -> CallableTy [ListTy t] (ListTy t) + BuiltinSum -> CallableTy [ListTy IntTy] IntTy + BuiltinTuple ts -> CallableTy [TupleTy ts] (TupleTy ts) + BuiltinZip ts -> CallableTy (map ListTy ts) (TupleTy ts) + BuiltinInput -> CallableTy [] StringTy + BuiltinPrint ts -> CallableTy ts SideEffectTy + +mapTypeBuiltin :: (Type -> Type) -> Builtin -> Builtin +mapTypeBuiltin f = \case + BuiltinAbs -> BuiltinAbs + BuiltinPow -> BuiltinPow + BuiltinModPow -> BuiltinModPow + BuiltinAll -> BuiltinAll + BuiltinAny -> BuiltinAny + BuiltinArgMax t -> BuiltinArgMax (f t) + BuiltinArgMin t -> BuiltinArgMin (f t) + BuiltinBool t -> BuiltinBool (f t) + BuiltinCeilDiv -> BuiltinCeilDiv + BuiltinCeilMod -> BuiltinCeilMod + BuiltinChoose -> BuiltinChoose + BuiltinDivMod -> BuiltinDivMod + BuiltinEnumerate t -> BuiltinEnumerate (f t) + BuiltinFact -> BuiltinFact + BuiltinFilter t -> BuiltinFilter (f t) + BuiltinFloorDiv -> BuiltinFloorDiv + BuiltinFloorMod -> BuiltinFloorMod + BuiltinGcd -> BuiltinGcd + BuiltinInt t -> BuiltinInt (f t) + BuiltinModInv -> BuiltinModInv + BuiltinLcm -> BuiltinLcm + BuiltinLen t -> BuiltinLen (f t) + BuiltinList t -> BuiltinList (f t) + BuiltinMap args ret -> BuiltinMap (map f args) (f ret) + BuiltinMax t n -> BuiltinMax (f t) n + BuiltinMax1 t -> BuiltinMax1 (f t) + BuiltinMin t n -> BuiltinMin (f t) n + BuiltinMin1 t -> BuiltinMin1 (f t) + BuiltinMultiChoose -> BuiltinMultiChoose + BuiltinPermute -> BuiltinPermute + BuiltinProduct -> BuiltinProduct + BuiltinRange1 -> BuiltinRange1 + BuiltinRange2 -> BuiltinRange2 + BuiltinRange3 -> BuiltinRange3 + BuiltinReversed t -> BuiltinReversed (f t) + BuiltinSorted t -> BuiltinSorted (f t) + BuiltinSum -> BuiltinSum + BuiltinTuple ts -> BuiltinTuple (map f ts) + BuiltinZip ts -> BuiltinZip (map f ts) + BuiltinInput -> BuiltinInput + BuiltinPrint ts -> BuiltinPrint (map f ts) + +attributeNames :: S.Set AttributeName +attributeNames = + S.fromList + [ "count", + "index", + "copy", + "append", + "split" + ] + +resolveAttribute' :: (MonadAlpha m, MonadError Error m) => Attribute' -> m Attribute' +resolveAttribute' x = wrapAt' (loc' x) $ case value' x of + UnresolvedAttribute x' -> + if x' `S.notMember` attributeNames + then throwSymbolError $ "unknown attribute: " ++ unAttributeName x' + else wrapError' "Jikka.RestrictedPython.Language.Builtin.resolveAttribute" $ do + WithLoc' (loc' x) <$> case x' of + "count" -> BuiltinCount <$> genType + "index" -> BuiltinIndex <$> genType + "copy" -> BuiltinCopy <$> genType + "append" -> BuiltinAppend <$> genType + "split" -> return BuiltinSplit + _ -> throwInternalError $ "not exhaustive: " ++ unAttributeName x' + _ -> return x + +resolveAttribute :: (MonadAlpha m, MonadError Error m) => Expr' -> Attribute' -> m Expr +resolveAttribute e@(WithLoc' _ (Name (WithLoc' _ "math"))) x = wrapAt' (loc' x) $ case value' x of + UnresolvedAttribute x' -> case x' of + "gcd" -> return (Constant (ConstBuiltin BuiltinGcd)) + "lcm" -> return (Constant (ConstBuiltin BuiltinGcd)) + _ -> throwSymbolError $ "unknown attribute: " ++ unAttributeName x' + _ -> return $ Attribute e x +resolveAttribute e@(WithLoc' _ (Name (WithLoc' _ "jikka"))) x = wrapAt' (loc' x) $ case value' x of + UnresolvedAttribute x' -> + let x'' = VarName (unAttributeName x') + in if x'' `S.notMember` additionalBuiltinNames + then throwSymbolError $ "unknown attribute: " ++ unAttributeName x' + else value' <$> resolveUniqueBuiltin (x $> x'') + _ -> return $ Attribute e x +resolveAttribute e x = Attribute e <$> resolveAttribute' x + +formatAttribute :: Attribute -> String +formatAttribute = \case + UnresolvedAttribute x -> unAttributeName x + BuiltinCount _ -> "count" + BuiltinIndex _ -> "index" + BuiltinCopy _ -> "copy" + BuiltinAppend _ -> "append" + BuiltinSplit -> "split" + +typeAttribute :: Attribute -> (Type, Type) +typeAttribute = \case + UnresolvedAttribute x -> error $ "Jikka.RestrictedPython.Language.Builtin.typeAttribute: attributes must be resolved: " ++ unAttributeName x + BuiltinCount t -> (ListTy t, CallableTy [t] IntTy) + BuiltinIndex t -> (ListTy t, CallableTy [t] IntTy) + BuiltinCopy t -> (ListTy t, CallableTy [] (ListTy t)) + BuiltinAppend t -> (ListTy t, CallableTy [t] SideEffectTy) + BuiltinSplit -> (StringTy, CallableTy [] (ListTy StringTy)) + +mapTypeAttribute :: (Type -> Type) -> Attribute -> Attribute +mapTypeAttribute f = \case + UnresolvedAttribute x -> UnresolvedAttribute x + BuiltinCount t -> BuiltinCount (f t) + BuiltinIndex t -> BuiltinIndex (f t) + BuiltinCopy t -> BuiltinCopy (f t) + BuiltinAppend t -> BuiltinAppend (f t) + BuiltinSplit -> BuiltinSplit diff --git a/src/Jikka/RestrictedPython/Language/Expr.hs b/src/Jikka/RestrictedPython/Language/Expr.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Language/Expr.hs @@ -0,0 +1,327 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE PatternSynonyms #-} + +-- | +-- Module : Jikka.RestrictedPython.Language.Expr +-- Description : contains data types of the restricted Python. / 制限された Python のためのデータ型を含みます。 +-- Copyright : (c) Kimiyuki Onaka, 2020 +-- License : Apache License 2.0 +-- Maintainer : kimiyuki95@gmail.com +-- Stability : experimental +-- Portability : portable +module Jikka.RestrictedPython.Language.Expr + ( -- * types + TypeName (..), + unTypeName, + Type (..), + pattern NoneTy, + + -- * operators + UnaryOp (..), + Operator (..), + BoolOp (..), + CmpOp (..), + CmpOp' (..), + Constant (..), + Builtin (..), + AttributeName (..), + unAttributeName, + Attribute (..), + Attribute', + + -- * exprs + VarName (..), + unVarName, + module Jikka.Common.Location, + VarName', + Expr (..), + Expr', + Comprehension (..), + + -- * statements + Target (..), + Target', + Statement (..), + pattern Append, + ToplevelStatement (..), + Program, + ) +where + +import Data.String (IsString) +import Jikka.Common.Location +import Jikka.Python.Language.Expr (BoolOp (..), CmpOp (..), Operator (..), UnaryOp (..)) + +newtype VarName = VarName String deriving (Eq, Ord, Show, Read, IsString) + +unVarName :: VarName -> String +unVarName (VarName x) = x + +type VarName' = WithLoc' VarName + +newtype TypeName = TypeName String deriving (Eq, Ord, Show, Read, IsString) + +unTypeName :: TypeName -> String +unTypeName (TypeName x) = x + +newtype AttributeName = AttributeName String deriving (Eq, Ord, Show, Read, IsString) + +unAttributeName :: AttributeName -> String +unAttributeName (AttributeName x) = x + +-- | `Type` represents the types of our restricted Python-like language. +-- +-- \[ +-- \newcommand\int{\mathbf{int}} +-- \newcommand\bool{\mathbf{bool}} +-- \newcommand\list{\mathbf{list}} +-- \newcommand\string{\mathbf{string}} +-- \begin{array}{rl} +-- \tau ::= & \alpha \\ +-- \vert & \int \\ +-- \vert & \bool \\ +-- \vert & \list(\tau) \\ +-- \vert & \tau \times \tau \times \dots \times \tau \\ +-- \vert & \tau \times \tau \times \dots \times \tau \to \tau +-- \vert & \string +-- \vert & \mathbf{side-effect} +-- \end{array} +-- \] +-- +-- NOTE: \(\mathbf{None}\) is represented as the 0-tuple. +data Type + = VarTy TypeName + | IntTy + | BoolTy + | ListTy Type + | TupleTy [Type] + | CallableTy [Type] Type + | StringTy + | SideEffectTy + deriving (Eq, Ord, Show, Read) + +pattern NoneTy = TupleTy [] + +data Constant + = ConstNone + | ConstInt Integer + | ConstBool Bool + | ConstBuiltin Builtin + deriving (Eq, Ord, Show, Read) + +data Builtin + = -- | "abs" \(: \int \to \int\) + BuiltinAbs + | -- | "pow" \((\lambda x k. x^k) : \int \times \int \to \int\) + BuiltinPow + | -- | modulo power "pow" \((\lambda x k m. x^k \bmod m): \int \times \int \to \int\) + BuiltinModPow + | -- | "divmod" \(: \int \times \int \to \int \times \int\) + BuiltinDivMod + | -- | ceil div \(: \int \times \int \to \int\) + BuiltinCeilDiv + | -- | ceil mod \(: \int \times \int \to \int\) + BuiltinCeilMod + | -- | floor div \(: \int \times \int \to \int\) + BuiltinFloorDiv + | -- | floor mod \(: \int \times \int \to \int\) + BuiltinFloorMod + | -- | \(\gcd: \int \times \int \to \int\) + BuiltinGcd + | -- | \(\mathbf{lcm}: \int \times \int \to \int\) + BuiltinLcm + | -- | "int" \(: \forall \alpha. \alpha \to \int\) + BuiltinInt Type + | -- | "bool" \(: \forall \alpha. \alpha \to \bool\) + BuiltinBool Type + | -- | "list" \(: \forall \alpha. \list(\alpha) \to \list(\alpha)\) + BuiltinList Type + | -- | "tuple" \(: \forall \alpha_0 \alpha_1 \dots \alpha _ {n - 1}. \tau \to \tau\) where \(\tau = \alpha_0 \times \dots \times \alpha _ {n - 1}\) + BuiltinTuple [Type] + | -- | "len" \(: \forall \alpha. \list(\alpha) \to \int\) + BuiltinLen Type + | -- | "map" \(: \forall \alpha_0 \alpha_1 \dots \alpha_n. (\alpha_0 \times \dots \times \alpha _ {n - 1} \to \alpha_n) \times \list(\alpha_0) \times \dots \list(\alpha _ {n - 1}) \to \list(\alpha_n)\) + BuiltinMap [Type] Type + | -- | "sorted" \(: \forall \alpha. \list(\alpha) \to \list(\alpha)\) + BuiltinSorted Type + | -- | "reversed" \(: \forall \alpha. \list(\alpha) \to \list(\alpha)\) + BuiltinReversed Type + | -- | "enumerate" \(: \forall \alpha. \list(\alpha) \to \list(\int \times \alpha)\) + BuiltinEnumerate Type + | -- | "filter" \(: \forall \alpha. (\alpha \to \bool) \times \list(\alpha) \to \list(\alpha)\) + BuiltinFilter Type + | -- | "zip" \(: \forall \alpha_0 \alpha_1 \dots \alpha _ {n - 1}. \list(\alpha_0) \times \dots \list(\alpha _ {n - 1}) \to \list(\alpha_0 \times \dots \times \alpha _ {n - 1})\) + BuiltinZip [Type] + | -- | "all" \(: \list(\bool) \to \bool\) + BuiltinAll + | -- | "any" \(: \list(\bool) \to \bool\) + BuiltinAny + | -- | "sum" \(: \list(\int) \to \int\) + BuiltinSum + | -- | product \(: \list(\int) \to \int\) + BuiltinProduct + | -- | "range" \(: \int \to \list(\int)\) + BuiltinRange1 + | -- | "range" \(: \int \times \int \to \list(\int)\) + BuiltinRange2 + | -- | "range" \(: \int \times \int \times \int \to \list(\int)\) + BuiltinRange3 + | -- | "max" \(: \forall \alpha. \list(\alpha) \to \alpha\) + BuiltinMax1 Type + | -- | "max" \(: \forall \alpha. \underbrace{\alpha \times \alpha \times \dots \times \alpha} _ {n ~\text{times}} \to \alpha\) + BuiltinMax Type Int + | -- | "min" \(: \forall \alpha. \list(\alpha) \to \alpha\) + BuiltinMin1 Type + | -- | "min" \(: \forall \alpha. \underbrace{\alpha \times \alpha \times \dots \times \alpha} _ {n ~\text{times}} \to \alpha\) + BuiltinMin Type Int + | -- | \(: \forall \alpha. \list(\alpha) \to \int\) + BuiltinArgMax Type + | -- | \(: \forall \alpha. \list(\alpha) \to \int\) + BuiltinArgMin Type + | -- | factorial \((\lambda n. n!): \int \to \int\) + BuiltinFact + | -- | \((\lambda n r. {} _ n C _ r): \int \times \int \to \int\) + BuiltinChoose + | -- | \((\lambda n r. {} _ n P _ r): \int \times \int \to \int\) + BuiltinPermute + | -- | \((\lambda n r. {} _ n H _ r): \int \times \int \to \int\) + BuiltinMultiChoose + | -- | modulo inverse \((\lambda x m. x^{-1} \bmod m): \int \times \int \to \int\) + BuiltinModInv + | -- | "input" \(: \epsilon \to \string\) + BuiltinInput + | -- | "print" \(: \forall \alpha_0 \alpha_1 \dots \alpha _ {n - 1}. \alpha_0 \times \dots \alpha _ {n - 1} \to \epsilon\) + BuiltinPrint [Type] + deriving (Eq, Ord, Show, Read) + +data Attribute + = UnresolvedAttribute AttributeName + | -- | "list.count" \(: \forall \alpha. \list(\alpha) \to \alpha \to \int\) + BuiltinCount Type + | -- | "list.index" \(: \forall \alpha. \list(\alpha) \to \alpha \to \int\) + BuiltinIndex Type + | -- | "list.copy" \(: \forall \alpha. \list(\alpha) \to \epsilon \to \list(\alpha)\) + BuiltinCopy Type + | -- | "list.append" \(: \forall \alpha. \list(\alpha) \to \alpha \to \mathbf{side-effect}\) + BuiltinAppend Type + | -- | "str.split" \(: \forall \alpha. \string \to \epsilon \to \list(\string)\) + BuiltinSplit + deriving (Eq, Ord, Show, Read) + +type Attribute' = WithLoc' Attribute + +-- | `Target` represents the lvalue of our restricted Python-like language. +-- +-- \[ +-- \begin{array}{rl} +-- y ::= & y \lbrack e \rbrack \\ +-- \vert & x \\ +-- \vert & (y, y, \dots, y) \\ +-- \end{array} +-- \] +data Target + = SubscriptTrg Target' Expr' + | NameTrg VarName' + | TupleTrg [Target'] + deriving (Eq, Ord, Show, Read) + +type Target' = WithLoc' Target + +-- | `CmpOp'` is a type for comparision operators. +-- This is annotated with its type as let-polymorphism. +data CmpOp' = CmpOp' CmpOp Type + deriving (Eq, Ord, Show, Read) + +data Comprehension = Comprehension Target' Expr' (Maybe Expr') + deriving (Eq, Ord, Show, Read) + +-- | `Expr` represents the exprs of our restricted Python-like language. +-- +-- \[ +-- \begin{array}{rl} +-- e ::= & e \operatorname{boolop} e \\ +-- \vert & e \operatorname{binop} e \\ +-- \vert & \operatorname{unaryop} e \\ +-- \vert & \lambda x _ \tau x _ \tau \dots x _ \tau. e \\ +-- \vert & \mathbf{if}~ e ~\mathbf{then}~ e ~\mathbf{else}~ e \\ +-- \vert & \lbrack e ~\mathbf{for}~ y ~\mathbf{in}~ e ~(\mathbf{if}~ e)? \rbrack \\ +-- \vert & e \operatorname{cmpop} e \\ +-- \vert & e (e, e, \dots, e) \\ +-- \vert & \operatorname{constant} \\ +-- \vert & e \lbrack e \rbrack \\ +-- \vert & x \\ +-- \vert & \lbrack e, e, \dots, e \rbrack _ \tau \\ +-- \vert & e \lbrack e? \colon e? \colon e? \rbrack \\ +-- \end{array} +-- \] +data Expr + = BoolOp Expr' BoolOp Expr' + | BinOp Expr' Operator Expr' + | UnaryOp UnaryOp Expr' + | Lambda [(VarName', Type)] Expr' + | IfExp Expr' Expr' Expr' + | ListComp Expr' Comprehension + | Compare Expr' CmpOp' Expr' + | Call Expr' [Expr'] + | Constant Constant + | Attribute Expr' Attribute' + | Subscript Expr' Expr' + | Starred Expr' + | Name VarName' + | List Type [Expr'] + | Tuple [Expr'] + | SubscriptSlice Expr' (Maybe Expr') (Maybe Expr') (Maybe Expr') + deriving (Eq, Ord, Show, Read) + +type Expr' = WithLoc' Expr + +-- | `Statement` represents the statements of our restricted Python-like language. +-- They appear in bodies of `def`. +-- +-- \[ +-- \begin{array}{rl} +-- \mathrm{stmt} ::= & \mathbf{return}~ e \\ +-- \vert & y \operatorname{binop} = e \\ +-- \vert & y _ \tau := e \\ +-- \vert & \mathbf{for}~ y ~\mathbf{in}~ e \colon\quad \mathrm{stmt}; \mathrm{stmt}; \dots; \mathrm{stmt} \\ +-- \vert & \mathbf{if}~ e \colon\quad \mathrm{stmt}; \mathrm{stmt}; \dots; \mathrm{stmt};\quad \mathbf{else}\colon\quad \mathrm{stmt}; \mathrm{stmt}; \dots; \mathrm{stmt} \\ +-- \vert & \mathbf{assert}~ e \\ +-- \end{array} +-- \] +data Statement + = Return Expr' + | AugAssign Target' Operator Expr' + | AnnAssign Target' Type Expr' + | For Target' Expr' [Statement] + | If Expr' [Statement] [Statement] + | Assert Expr' + | -- | expression statements + Expr' Expr' + deriving (Eq, Ord, Show, Read) + +pattern Append loc t e1 e2 <- Expr' (WithLoc' loc (Call (WithLoc' _ (Attribute e1 (WithLoc' _ (BuiltinAppend t)))) [e2])) + +-- | `TopLevelStatement` represents the statements of our restricted Python-like language. +-- They appear in the toplevel of programs. +-- +-- \[ +-- \begin{array}{rl} +-- \mathrm{tlstmt} ::= & x _ \tau := e \\ +-- \vert & \mathbf{def}~ x (x _ \tau, x _ \tau, \dots, x _ \tau) \to \tau \colon\quad \mathrm{stmt}; \mathrm{stmt}; \dots; \mathrm{stmt} \\ +-- \vert & \mathbf{assert}~ e \\ +-- \end{array} +-- \] +data ToplevelStatement + = ToplevelAnnAssign VarName' Type Expr' + | ToplevelFunctionDef VarName' [(VarName', Type)] Type [Statement] + | ToplevelAssert Expr' + deriving (Eq, Ord, Show, Read) + +-- | `Program` represents the programs of our restricted Python-like language. +-- +-- \[ +-- \begin{array}{rl} +-- \mathrm{prog} ::= & \mathrm{tlstmt}; \mathrm{tlstmt}; \dots; \mathrm{tlstmt} \\ +-- \end{array} +-- \] +type Program = [ToplevelStatement] diff --git a/src/Jikka/RestrictedPython/Language/Lint.hs b/src/Jikka/RestrictedPython/Language/Lint.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Language/Lint.hs @@ -0,0 +1,264 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +module Jikka.RestrictedPython.Language.Lint where + +import Control.Monad.Writer.Strict +import qualified Data.Set as S +import Jikka.Common.Error +import Jikka.RestrictedPython.Language.Builtin (builtinNames) +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Util +import Jikka.RestrictedPython.Language.VariableAnalysis + +makeEnsureProgram :: MonadError Error m => (Program -> Bool) -> String -> Program -> m () +makeEnsureProgram pred msg prog = + unless (pred prog) $ do + throwSemanticError msg + +-- | `hasSubscriptionInLoopCounters` checks that there are `SubscriptTrg` in loop counters of for-loops. +-- This includes loop counters of `ListComp`. +-- For example, the followings has such subscriptions. +-- +-- > for a[0] in range(100): +-- > pass +-- > return a[0] # => 99 +-- +-- > a = [0] +-- > b = [0 for a[0] in range(100)] +-- > return a[0] # => 99 +-- +-- NOTE: This is allowd in the standard Python. +hasSubscriptionInLoopCounters :: Program -> Bool +hasSubscriptionInLoopCounters prog = any checkStatement (listStatements prog) || any checkExpr (listExprs prog) + where + checkStatement = \case + For x _ _ -> hasSubscriptTrg x + _ -> False + checkExpr (WithLoc' _ x) = case x of + ListComp _ (Comprehension x _ _) -> hasSubscriptTrg x + _ -> False + +doesntHaveSubscriptionInLoopCounters :: Program -> Bool +doesntHaveSubscriptionInLoopCounters = not . hasSubscriptionInLoopCounters + +ensureDoesntHaveSubscriptionInLoopCounters :: MonadError Error m => Program -> m () +ensureDoesntHaveSubscriptionInLoopCounters = makeEnsureProgram doesntHaveSubscriptionInLoopCounters "there must not be subscription in loop counters" + +-- | `hasLeakOfLoopCounters` checks that there are leaks of loop counters of for-loops. +-- For example, the following has a leak. +-- +-- > for i in range(100): +-- > pass +-- > return i # => 100 +hasLeakOfLoopCounters :: Program -> Bool +hasLeakOfLoopCounters _ = False -- TODO + +doesntHaveLeakOfLoopCounters :: Program -> Bool +doesntHaveLeakOfLoopCounters = not . hasLeakOfLoopCounters + +ensureDoesntHaveLeakOfLoopCounters :: MonadError Error m => Program -> m () +ensureDoesntHaveLeakOfLoopCounters = makeEnsureProgram doesntHaveLeakOfLoopCounters "there must not be leaks of loop counters" + +-- | `hasAssignmentToLoopCounters` checks that there are assignments to loop counters of for-loops. +-- For example, the following has the assignment. +-- +-- > for i in range(100): +-- > i += 1 +hasAssignmentToLoopCounters :: Program -> Bool +hasAssignmentToLoopCounters prog = any check (listStatements prog) + where + check = \case + For x _ body -> + let r = ReadList $ targetVars x + (_, w) = analyzeStatementsMax body + in haveWriteReadIntersection w r + _ -> False + +doesntHaveAssignmentToLoopCounters :: Program -> Bool +doesntHaveAssignmentToLoopCounters = not . hasAssignmentToLoopCounters + +ensureDoesntHaveAssignmentToLoopCounters :: MonadError Error m => Program -> m () +ensureDoesntHaveAssignmentToLoopCounters = makeEnsureProgram doesntHaveAssignmentToLoopCounters "there must not be assignments to loop counters" + +-- | `hasAssignmentToLoopIterators` checks that there are assignments to loop iterators of for-loops. +-- For example, the followings have the assignments. +-- +-- > a = list(range(10)) +-- > for i in a: +-- > a[5] = i +-- +-- > a = 0 +-- > for i in f(a): +-- > a += i +hasAssignmentToLoopIterators :: Program -> Bool +hasAssignmentToLoopIterators prog = any check (listStatements prog) + where + check = \case + For _ iter body -> + let r = analyzeExpr iter + (_, w) = analyzeStatementsMax body + in haveWriteReadIntersection w r + _ -> False + +doesntHaveAssignmentToLoopIterators :: Program -> Bool +doesntHaveAssignmentToLoopIterators = not . hasAssignmentToLoopIterators + +ensureDoesntHaveAssignmentToLoopIterators :: MonadError Error m => Program -> m () +ensureDoesntHaveAssignmentToLoopIterators = makeEnsureProgram doesntHaveAssignmentToLoopIterators "there must not be assignments changing loop iterators" + +-- | `hasReturnInLoops` checks that there are return-statements in for-loops. +-- For example, the following has such a return-statement. +-- +-- > a = list(range(10)) +-- > for i in a: +-- > return True +hasReturnInLoops :: Program -> Bool +hasReturnInLoops = getAny . execWriter . mapLargeStatementM fIf fFor + where + fIf e body1 body2 = return [If e body1 body2] + fFor x iter body = do + when (any doesPossiblyReturn body) $ do + tell $ Any True + return [For x iter body] + +doesntHaveReturnInLoops :: Program -> Bool +doesntHaveReturnInLoops = not . hasReturnInLoops + +ensureDoesntHaveReturnInLoops :: MonadError Error m => Program -> m () +ensureDoesntHaveReturnInLoops = makeEnsureProgram doesntHaveReturnInLoops "there must not be return-statements in for-loops" + +-- | `hasMixedAssignment` checks that there are assignments which assign to both of bare variables and subscripted variables. +-- For example, the following is such an assignment. +-- +-- > a, b[0] = list(range(10)) +-- +-- NOTE: this doesn't check loop counters of `For` or `ListComp`. +hasMixedAssignment :: Program -> Bool +hasMixedAssignment prog = any check (listStatements prog) + where + check = \case + AugAssign x _ _ -> hasSubscriptTrg x && hasBareNameTrg x + AnnAssign x _ _ -> hasSubscriptTrg x && hasBareNameTrg x + _ -> False + +doesntHaveMixedAssignment :: Program -> Bool +doesntHaveMixedAssignment = not . hasMixedAssignment + +ensureDoesntHaveMixedAssignment :: MonadError Error m => Program -> m () +ensureDoesntHaveMixedAssignment = makeEnsureProgram doesntHaveMixedAssignment "there must not be mixed assignments" + +-- | `hasNonTrivialSubscriptedAssignmentInForLoops` checks that there are assignments with non-trivial subscriptions in for-loops. +-- A trivial subscription is a sequence of subscriptions to a variable with constant indices and at most one trivial loop-counter indices for each loops. +-- A constant index is an expr which has a constant value in the loop. +-- A trivial loop-counter index is the loop counter from "range(n)", "range(n, m)" or "enumerate(a)" with optional post-addition with a positive int literal. +-- +-- For example, the followings have such assignments. +-- +-- > x = 0 +-- > for i in range(10): +-- > x += 1 +-- > a[x] += 1 +-- +-- > for i in range(10): +-- > j = i +-- > a[j] += 1 +-- +-- > for i in range(10): +-- > a[2 * i] += 1 +-- +-- > for i in range(10): +-- > a[1 + i] += 1 +-- +-- > for i in range(10): +-- > a[i - 1] += 1 +-- +-- > c = 1 +-- > for i in range(10): +-- > a[i + c] += 1 +-- +-- > for i in range(10): +-- > a[i][i] += 1 +-- +-- > for i in [1, 2, 3]: +-- > a[i] += 1 +-- +-- > b = range(10) +-- > for i in b: +-- > a[i] += 1 +-- +-- > for i in range(0, 10, 2): +-- > a[i] += 1 +-- +-- > for i, b_i in enumerate(b): +-- > a[b_i] += i +-- +-- For example, the followings don't have such assignments. +-- +-- > c = 0 +-- > for i in range(10): +-- > a[c] += 1 +-- +-- > for i in range(10): +-- > a[i] += 1 +-- +-- > for i in range(10): +-- > a[i + 1] += 1 +-- +-- > for i in range(10): +-- > for j in range(10): +-- > a[i + 1][j] += 1 +-- +-- > for i in range(1, 10): +-- > a[i] += 1 +-- +-- > for i, b_i in enumerate(b): +-- > a[i] += b_i +hasNonTrivialSubscriptedAssignmentInForLoops :: Program -> Bool +hasNonTrivialSubscriptedAssignmentInForLoops prog = any check (listStatements prog) + where + check = \case + AugAssign x _ _ -> go x + AnnAssign x _ _ -> go x + _ -> False + go (WithLoc' _ x) = case x of + SubscriptTrg _ _ -> False -- TODO + NameTrg _ -> False + TupleTrg xs -> any go xs + +doesntHaveNonTrivialSubscriptedAssignmentInForLoops :: Program -> Bool +doesntHaveNonTrivialSubscriptedAssignmentInForLoops = not . hasMixedAssignment + +ensureDoesntHaveNonTrivialSubscriptedAssignmentInForLoops :: MonadError Error m => Program -> m () +ensureDoesntHaveNonTrivialSubscriptedAssignmentInForLoops = makeEnsureProgram doesntHaveNonTrivialSubscriptedAssignmentInForLoops "there must not be assignments with non-trivial subscriptions in for-loops" + +-- | `hasAssginmentToBuiltin` checks that there are assignments to builtin functions. +-- For example, the followings have such assignments. +-- +-- > map = 3 +-- +-- > return [range for range in range(10)] +hasAssignmentToBuiltin :: Program -> Bool +hasAssignmentToBuiltin _ = False -- TODO + +doesntHaveAssignmentToBuiltin :: Program -> Bool +doesntHaveAssignmentToBuiltin = not . hasAssignmentToBuiltin + +ensureDoesntHaveAssignmentToBuiltin :: MonadError Error m => Program -> m () +ensureDoesntHaveAssignmentToBuiltin = makeEnsureProgram doesntHaveAssignmentToBuiltin "there must not be assignments to builtin functions" + +-- | `hasNonResolvedBuiltin` checks that there are not resolved builtin functions. +-- This always doesn't hold after `Jikka.RestrictedPython.Language.Convert.ResolveBuiltin`. +hasNonResolvedBuiltin :: Program -> Bool +hasNonResolvedBuiltin = any check . listExprs + where + check = any check' . listSubExprs + check' (WithLoc' _ e) = case e of + Name x | value' x `S.member` builtinNames -> True + _ -> False + +doesntHaveNonResolvedBuiltin :: Program -> Bool +doesntHaveNonResolvedBuiltin = not . hasAssignmentToBuiltin + +ensureDoesntHaveNonResolvedBuiltin :: MonadError Error m => Program -> m () +ensureDoesntHaveNonResolvedBuiltin = makeEnsureProgram doesntHaveNonResolvedBuiltin "there must not be assignments to builtin functions" diff --git a/src/Jikka/RestrictedPython/Language/Util.hs b/src/Jikka/RestrictedPython/Language/Util.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Language/Util.hs @@ -0,0 +1,356 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} + +module Jikka.RestrictedPython.Language.Util + ( -- * generating symbols + genType, + genVarName, + genVarName', + + -- * free variables + freeTyVars, + freeVars, + freeVars', + freeVarsTarget, + freeVarsTarget', + + -- * return-statements + doesAlwaysReturn, + doesPossiblyReturn, + + -- * traversing statements + mapStatement, + mapStatementM, + mapLargeStatement, + mapLargeStatementM, + mapStatements, + mapStatementsM, + listStatements, + + -- * traversing sub exprs + mapSubExprM, + mapSubExpr, + listSubExprs, + + -- * traversing exprs + mapExprTargetM, + mapExprStatementM, + mapExprM, + listExprs, + + -- * exprs + hasFunctionCall, + isSmallExpr, + dropLocation, + + -- * targets + targetVars, + targetVars', + hasSubscriptTrg, + hasBareNameTrg, + exprToTarget, + targetToExpr, + + -- * programs + toplevelMainDef, + ) +where + +import Control.Monad.Identity +import Control.Monad.Writer.Strict +import Data.List (delete, nub) +import Jikka.Common.Alpha +import Jikka.Common.Location +import Jikka.RestrictedPython.Language.Expr + +genType :: MonadAlpha m => m Type +genType = do + i <- nextCounter + return $ VarTy (TypeName ('$' : show i)) + +genVarName :: MonadAlpha m => VarName' -> m VarName' +genVarName x = do + i <- nextCounter + let base = if unVarName (value' x) == "_" then "" else takeWhile (/= '$') (unVarName (value' x)) + return $ WithLoc' (loc' x) (VarName (base ++ '$' : show i)) + +genVarName' :: MonadAlpha m => m VarName' +genVarName' = genVarName (withoutLoc (VarName "_")) + +freeTyVars :: Type -> [TypeName] +freeTyVars = nub . go + where + go = \case + VarTy x -> [x] + IntTy -> [] + BoolTy -> [] + ListTy t -> go t + TupleTy ts -> concat $ mapM go ts + CallableTy ts ret -> concat $ mapM go (ret : ts) + StringTy -> [] + SideEffectTy -> [] + +-- | `freeVars'` reports all free variables. +freeVars :: Expr' -> [VarName] +freeVars = nub . map value' . freeVars' + +-- | `freeVars'` reports all free variables with their locations, i.e. occurrences. +-- For examples, @x + x@ and @x@ have the same free variables @x@ but they have different sets of occurrences of free variable. +freeVars' :: Expr' -> [VarName'] +freeVars' (WithLoc' _ e0) = case e0 of + BoolOp e1 _ e2 -> freeVars' e1 ++ freeVars' e2 + BinOp e1 _ e2 -> freeVars' e1 ++ freeVars' e2 + UnaryOp _ e -> freeVars' e + Lambda args e -> foldl (\vars (x, _) -> delete x vars) (freeVars' e) args + IfExp e1 e2 e3 -> freeVars' e1 ++ freeVars' e2 ++ freeVars' e3 + ListComp e (Comprehension x iter pred) -> freeVars' iter ++ foldl (\vars x -> delete x vars) (freeVars' e ++ concatMap freeVars' pred) (targetVars' x) + Compare e1 _ e2 -> freeVars' e1 ++ freeVars' e2 + Call f args -> concatMap freeVars' (f : args) + Constant _ -> [] + Attribute e _ -> freeVars' e + Subscript e1 e2 -> freeVars' e1 ++ freeVars' e2 + Starred e -> freeVars' e + Name x -> [x] + List _ es -> concatMap freeVars' es + Tuple es -> concatMap freeVars' es + SubscriptSlice e from to step -> freeVars' e ++ concatMap freeVars' from ++ concatMap freeVars' to ++ concatMap freeVars' step + +freeVarsTarget :: Target' -> [VarName] +freeVarsTarget = nub . map value' . freeVarsTarget' + +freeVarsTarget' :: Target' -> [VarName'] +freeVarsTarget' (WithLoc' _ x) = case x of + SubscriptTrg _ e -> freeVars' e + NameTrg _ -> [] + TupleTrg xs -> concatMap freeVarsTarget' xs + +doesAlwaysReturn :: Statement -> Bool +doesAlwaysReturn = \case + Return _ -> True + AugAssign _ _ _ -> False + AnnAssign _ _ _ -> False + For _ _ _ -> False + If _ body1 body2 -> any doesAlwaysReturn body1 && any doesAlwaysReturn body2 + Assert _ -> False + Expr' _ -> False + +doesPossiblyReturn :: Statement -> Bool +doesPossiblyReturn = \case + Return _ -> True + AugAssign _ _ _ -> False + AnnAssign _ _ _ -> False + For _ _ body -> any doesPossiblyReturn body + If _ body1 body2 -> any doesPossiblyReturn body1 || any doesPossiblyReturn body2 + Assert _ -> False + Expr' _ -> False + +-- | `mapSubExprM` replaces all exprs in a given expr using a given function. +-- This may breaks various constraints. +mapSubExprM :: Monad m => (Expr' -> m Expr') -> Expr' -> m Expr' +mapSubExprM f = go + where + go e0 = + f . WithLoc' (loc' e0) =<< case value' e0 of + BoolOp e1 op e2 -> BoolOp <$> go e1 <*> return op <*> go e2 + BinOp e1 op e2 -> BinOp <$> go e1 <*> return op <*> go e2 + UnaryOp op e -> UnaryOp op <$> go e + Lambda args body -> Lambda args <$> go body + IfExp e1 e2 e3 -> IfExp <$> go e1 <*> go e2 <*> go e3 + ListComp e (Comprehension x iter pred) -> do + e <- go e + x <- mapExprTargetM f x + iter <- go iter + pred <- mapM go pred + return $ ListComp e (Comprehension x iter pred) + Compare e1 op e2 -> Compare <$> go e1 <*> return op <*> go e2 + Call g args -> Call <$> go g <*> mapM go args + Constant const -> return $ Constant const + Attribute e x -> Attribute <$> go e <*> pure x + Subscript e1 e2 -> Subscript <$> go e1 <*> go e2 + Starred e -> Starred <$> go e + Name x -> return $ Name x + List t es -> List t <$> mapM go es + Tuple es -> Tuple <$> mapM go es + SubscriptSlice e from to step -> SubscriptSlice <$> go e <*> mapM go from <*> mapM go to <*> mapM go step + +mapSubExpr :: (Expr' -> Expr') -> Expr' -> Expr' +mapSubExpr f = runIdentity . mapSubExprM (return . f) + +listSubExprs :: Expr' -> [Expr'] +listSubExprs = reverse . getDual . execWriter . mapSubExprM go + where + go e = do + tell $ Dual [e] + return e + +mapExprTargetM :: Monad m => (Expr' -> m Expr') -> Target' -> m Target' +mapExprTargetM f x = + WithLoc' (loc' x) <$> case value' x of + SubscriptTrg x e -> SubscriptTrg <$> mapExprTargetM f x <*> f e + NameTrg x -> return $ NameTrg x + TupleTrg xs -> TupleTrg <$> mapM (mapExprTargetM f) xs + +mapExprStatementM :: Monad m => (Expr' -> m Expr') -> Statement -> m Statement +mapExprStatementM f = \case + Return e -> Return <$> f e + AugAssign x op e -> AugAssign <$> mapExprTargetM f x <*> pure op <*> f e + AnnAssign x t e -> AnnAssign <$> mapExprTargetM f x <*> pure t <*> f e + For x iter body -> For <$> mapExprTargetM f x <*> f iter <*> mapM (mapExprStatementM f) body + If e body1 body2 -> If <$> f e <*> mapM (mapExprStatementM f) body1 <*> mapM (mapExprStatementM f) body2 + Assert e -> Assert <$> f e + Expr' e -> Expr' <$> f e + +mapExprToplevelStatementM :: Monad m => (Expr' -> m Expr') -> ToplevelStatement -> m ToplevelStatement +mapExprToplevelStatementM f = \case + ToplevelAnnAssign x t e -> ToplevelAnnAssign x t <$> f e + ToplevelFunctionDef g args ret body -> ToplevelFunctionDef g args ret <$> mapM (mapExprStatementM f) body + ToplevelAssert e -> ToplevelAssert <$> f e + +mapExprM :: Monad m => (Expr' -> m Expr') -> Program -> m Program +mapExprM f = mapM (mapExprToplevelStatementM f) + +listExprs :: Program -> [Expr'] +listExprs = reverse . getDual . execWriter . mapExprM go + where + go e = do + tell $ Dual [e] + return e + +mapStatementStatementM :: Monad m => (Statement -> m [Statement]) -> Statement -> m [Statement] +mapStatementStatementM f = \case + Return e -> f $ Return e + AugAssign x op e -> f $ AugAssign x op e + AnnAssign x t e -> f $ AnnAssign x t e + For x iter body -> do + body <- concat <$> mapM (mapStatementStatementM f) body + f $ For x iter body + If e body1 body2 -> do + body1 <- concat <$> mapM (mapStatementStatementM f) body1 + body2 <- concat <$> mapM (mapStatementStatementM f) body2 + f $ If e body1 body2 + Assert e -> f $ Assert e + Expr' e -> f $ Expr' e + +mapStatementToplevelStatementM :: Monad m => (Statement -> m [Statement]) -> ToplevelStatement -> m ToplevelStatement +mapStatementToplevelStatementM go = \case + ToplevelAnnAssign x t e -> return $ ToplevelAnnAssign x t e + ToplevelFunctionDef f args ret body -> do + body <- concat <$> mapM (mapStatementStatementM go) body + return $ ToplevelFunctionDef f args ret body + ToplevelAssert e -> return $ ToplevelAssert e + +-- | `mapStatementM` replaces all statements in a given program using a given function. +-- This may breaks various constraints. +mapStatementM :: Monad m => (Statement -> m [Statement]) -> Program -> m Program +mapStatementM f = mapM (mapStatementToplevelStatementM f) + +mapStatement :: (Statement -> [Statement]) -> Program -> Program +mapStatement f = runIdentity . mapStatementM (return . f) + +mapLargeStatementM :: Monad m => (Expr' -> [Statement] -> [Statement] -> m [Statement]) -> (Target' -> Expr' -> [Statement] -> m [Statement]) -> Program -> m Program +mapLargeStatementM fIf fFor = mapStatementM go + where + go = \case + Return e -> return [Return e] + AugAssign x op e -> return [AugAssign x op e] + AnnAssign x t e -> return [AnnAssign x t e] + For x iter body -> fFor x iter body + If e body1 body2 -> fIf e body1 body2 + Assert e -> return [Assert e] + Expr' e -> return [Expr' e] + +mapLargeStatement :: (Expr' -> [Statement] -> [Statement] -> [Statement]) -> (Target' -> Expr' -> [Statement] -> [Statement]) -> Program -> Program +mapLargeStatement fIf fFor = runIdentity . mapLargeStatementM fIf' fFor' + where + fIf' e body1 body2 = return $ fIf e body1 body2 + fFor' x iter body = return $ fFor x iter body + +listStatements :: Program -> [Statement] +listStatements = reverse . getDual . execWriter . mapStatementM go + where + go stmt = do + tell $ Dual [stmt] + return [stmt] + +mapStatementsToplevelStatementM :: Monad m => ([Statement] -> m [Statement]) -> ToplevelStatement -> m ToplevelStatement +mapStatementsToplevelStatementM go = \case + ToplevelAnnAssign x t e -> return $ ToplevelAnnAssign x t e + ToplevelFunctionDef f args ret body -> do + let go' = \case + Return e -> return [Return e] + AugAssign x op e -> return [AugAssign x op e] + AnnAssign x t e -> return [AnnAssign x t e] + For x iter body -> do + body <- go body + return [For x iter body] + If e body1 body2 -> do + body1 <- go body1 + body2 <- go body2 + return [If e body1 body2] + Assert e -> return [Assert e] + Expr' e -> return [Expr' e] + body <- concat <$> mapM (mapStatementStatementM go') body + body <- go body + return $ ToplevelFunctionDef f args ret body + ToplevelAssert e -> return $ ToplevelAssert e + +mapStatementsM :: Monad m => ([Statement] -> m [Statement]) -> Program -> m Program +mapStatementsM f = mapM (mapStatementsToplevelStatementM f) + +mapStatements :: ([Statement] -> [Statement]) -> Program -> Program +mapStatements f = runIdentity . mapStatementsM (return . f) + +hasFunctionCall :: Expr' -> Bool +hasFunctionCall = any (check . value') . listSubExprs + where + check = \case + Call _ _ -> True + _ -> False + +-- | `isSmallExpr` is true if the evaluation of a given expr trivially terminates. +isSmallExpr :: Expr' -> Bool +isSmallExpr = not . hasFunctionCall + +dropLocation :: Expr' -> Expr' +dropLocation = mapSubExpr go + where + go (WithLoc' _ e) = withoutLoc e + +targetVars :: Target' -> [VarName] +targetVars = nub . map value' . targetVars' + +targetVars' :: Target' -> [VarName'] +targetVars' (WithLoc' _ x) = case x of + SubscriptTrg x _ -> targetVars' x + NameTrg x -> [x] + TupleTrg xs -> concatMap targetVars' xs + +hasSubscriptTrg :: Target' -> Bool +hasSubscriptTrg (WithLoc' _ x) = case x of + SubscriptTrg _ _ -> True + NameTrg _ -> False + TupleTrg xs -> any hasSubscriptTrg xs + +hasBareNameTrg :: Target' -> Bool +hasBareNameTrg (WithLoc' _ x) = case x of + SubscriptTrg _ _ -> False + NameTrg _ -> True + TupleTrg xs -> any hasSubscriptTrg xs + +exprToTarget :: Expr' -> Maybe Target' +exprToTarget e = + WithLoc' (loc' e) <$> case value' e of + Name x -> Just $ NameTrg x + Tuple es -> TupleTrg <$> mapM exprToTarget es + Subscript e1 e2 -> SubscriptTrg <$> exprToTarget e1 <*> pure e2 + _ -> Nothing + +targetToExpr :: Target' -> Expr' +targetToExpr e = + WithLoc' (loc' e) $ case value' e of + NameTrg x -> Name x + TupleTrg es -> Tuple (map targetToExpr es) + SubscriptTrg e1 e2 -> Subscript (targetToExpr e1) e2 + +toplevelMainDef :: [Statement] -> Program +toplevelMainDef body = [ToplevelFunctionDef (WithLoc' Nothing (VarName "main")) [] IntTy body] diff --git a/src/Jikka/RestrictedPython/Language/Value.hs b/src/Jikka/RestrictedPython/Language/Value.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Language/Value.hs @@ -0,0 +1,113 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.RestrictedPython.Language.Value where + +import Data.Char (toLower) +import Data.List (intercalate) +import qualified Data.Map.Strict as M +import Data.Maybe (fromMaybe) +import qualified Data.Vector as V +import Jikka.Common.Error +import Jikka.Common.IOFormat +import Jikka.Common.Matrix +import Jikka.RestrictedPython.Language.Expr + +-- | `Value` is the values of our restricted Python-like language. +-- +-- \[ +-- \begin{array}{rl} +-- v ::= & \dots, -2, -1, 0, 1, 2, \dots \\ +-- \vert & \mathbf{false}, \mathbf{true} \\ +-- \vert & \mathbf{nil} \\ +-- \vert & \mathbf{cons}(v, v) \\ +-- \vert & (v, v, \dots, v) \\ +-- \vert & \lambda _ \mu x x \dots x. e \\ +-- \vert & \mathrm{builtin} \\ +-- \end{array} +-- \] +data Value + = IntVal Integer + | BoolVal Bool + | ListVal (V.Vector Value) + | TupleVal [Value] + | ClosureVal Local [(VarName, Type)] [Statement] + | BuiltinVal Builtin + | AttributeVal Value Attribute + deriving (Eq, Ord, Show, Read) + +newtype Local = Local + { unLocal :: M.Map VarName Value + } + deriving (Eq, Ord, Show, Read) + +toInt :: MonadError Error m => Value -> m Integer +toInt = \case + IntVal n -> return n + v -> throwInternalError $ "not an integer value: " ++ formatValue v + +toBool :: MonadError Error m => Value -> m Bool +toBool = \case + BoolVal p -> return p + v -> throwInternalError $ "not a boolean value: " ++ formatValue v + +toList :: MonadError Error m => Value -> m (V.Vector Value) +toList = \case + ListVal xs -> return xs + v -> throwInternalError $ "not a list value: " ++ formatValue v + +toTuple :: MonadError Error m => Value -> m [Value] +toTuple = \case + TupleVal xs -> return xs + v -> throwInternalError $ "not a tuple value: " ++ formatValue v + +toIntList :: MonadError Error m => Value -> m (V.Vector Integer) +toIntList xs = V.mapM toInt =<< toList xs + +toBoolList :: MonadError Error m => Value -> m (V.Vector Bool) +toBoolList xs = V.mapM toBool =<< toList xs + +toMatrix :: MonadError Error m => Value -> m (Matrix Integer) +toMatrix a = toMatrix' =<< V.mapM toIntList =<< toList a + where + toMatrix' a = case makeMatrix a of + Just a -> return a + Nothing -> throwInternalError $ "not a matrix: " ++ show a + +fromMatrix :: Matrix Integer -> Value +fromMatrix a = ListVal (fmap (ListVal . fmap IntVal) (unMatrix a)) + +compareValues :: Value -> Value -> Maybe Ordering +compareValues a b = case (a, b) of + (IntVal a, IntVal b) -> Just $ compare a b + (BoolVal a, BoolVal b) -> Just $ compare a b + (ListVal a, ListVal b) -> case mconcat (V.toList (V.zipWith compareValues a b)) of + Nothing -> Nothing + Just EQ -> Just $ compare (V.length a) (V.length b) + Just o -> Just o + (TupleVal a, TupleVal b) -> + if length a /= length b + then Nothing + else mconcat (zipWith compareValues a b) + (_, _) -> Nothing + +compareValues' :: Value -> Value -> Ordering +compareValues' a b = fromMaybe EQ (compareValues a b) + +formatValue :: Value -> String +formatValue = \case + IntVal n -> show n + BoolVal p -> map toLower (show p) + ListVal xs -> "[" ++ intercalate ", " (map formatValue (V.toList xs)) ++ "]" + TupleVal [x] -> "(" ++ formatValue x ++ ",)" + TupleVal xs -> "(" ++ intercalate ", " (map formatValue xs) ++ ")" + f@ClosureVal {} -> show f + BuiltinVal b -> show b + AttributeVal x a -> "(" ++ formatValue x ++ ")." ++ show a + +readValueIO :: (MonadIO m, MonadError Error m) => IOFormat -> m ([Value], M.Map String Value) +readValueIO = makeReadValueIO toInt IntVal toList ListVal + +writeValueIO :: (MonadError Error m, MonadIO m) => IOFormat -> M.Map String Value -> Value -> m () +writeValueIO = makeWriteValueIO toTuple IntVal toInt toList diff --git a/src/Jikka/RestrictedPython/Language/VariableAnalysis.hs b/src/Jikka/RestrictedPython/Language/VariableAnalysis.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Language/VariableAnalysis.hs @@ -0,0 +1,82 @@ +{-# LANGUAGE LambdaCase #-} + +module Jikka.RestrictedPython.Language.VariableAnalysis where + +import Data.List (delete, intersect, nub) +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Util + +newtype ReadList = ReadList [VarName] + deriving (Eq, Ord, Show, Read) + +newtype WriteList = WriteList [VarName] + deriving (Eq, Ord, Show, Read) + +haveWriteReadIntersection :: WriteList -> ReadList -> Bool +haveWriteReadIntersection (WriteList w) (ReadList r) = not (null (w `intersect` r)) + +analyzeExpr :: Expr' -> ReadList +analyzeExpr = ReadList . freeVars + +analyzeTargetRead :: Target' -> ReadList +analyzeTargetRead = ReadList . freeVarsTarget + +analyzeTargetWrite :: Target' -> WriteList +analyzeTargetWrite = WriteList . targetVars + +analyzeStatementGeneric :: Bool -> Statement -> (ReadList, WriteList) +analyzeStatementGeneric isMax = \case + Return e -> (analyzeExpr e, WriteList []) + AugAssign x _ e -> + let w = analyzeTargetWrite x + (ReadList r) = analyzeTargetRead x + (ReadList r') = analyzeExpr e + in (ReadList (nub $ r ++ r'), w) + AnnAssign x _ e -> + let w = analyzeTargetWrite x + (ReadList r) = analyzeTargetRead x + (ReadList r') = analyzeExpr e + in (ReadList (nub $ r ++ r'), w) + For x iter body -> + let xs = targetVars x + ReadList r = analyzeExpr iter + (ReadList r', WriteList w) = analyzeStatementsGeneric isMax body + in if isMax + then (ReadList (nub $ r ++ foldl (flip delete) r' xs), WriteList (nub $ foldl (flip delete) w xs)) + else (ReadList r, WriteList []) + If e body1 body2 -> + let ReadList r = analyzeExpr e + (ReadList r1, WriteList w1) = analyzeStatementsGeneric isMax body1 + (ReadList r2, WriteList w2) = analyzeStatementsGeneric isMax body2 + in if isMax + then (ReadList (nub $ r ++ r1 ++ r2), WriteList (nub $ w1 ++ w2)) + else (ReadList (nub $ r ++ intersect r1 r2), WriteList (nub $ w1 `intersect` w2)) + Assert e -> (analyzeExpr e, WriteList []) + Append _ _ x e -> + let w = maybe (WriteList []) analyzeTargetWrite (exprToTarget x) + (ReadList r) = maybe (ReadList []) analyzeTargetRead (exprToTarget x) + (ReadList r') = analyzeExpr e + in (ReadList (nub $ r ++ r'), w) + Expr' e -> (analyzeExpr e, WriteList []) + +analyzeStatementsGeneric :: Bool -> [Statement] -> (ReadList, WriteList) +analyzeStatementsGeneric isMax = go [] [] + where + go r w [] = (ReadList (nub r), WriteList (nub w)) + go r w (stmt : stmts) = + let (ReadList r', WriteList w') = analyzeStatementGeneric isMax stmt + in go (r' ++ r) (w' ++ w) stmts + +-- | `analyzeStatementMax` returns lists of variables which are possibly read or written in given statements. +analyzeStatementMax :: Statement -> (ReadList, WriteList) +analyzeStatementMax = analyzeStatementGeneric True + +analyzeStatementsMax :: [Statement] -> (ReadList, WriteList) +analyzeStatementsMax = analyzeStatementsGeneric True + +-- | `analyzeStatementMin` returns lists of variables which are always read or written in given statements. +analyzeStatementMin :: Statement -> (ReadList, WriteList) +analyzeStatementMin = analyzeStatementGeneric False + +analyzeStatementsMin :: [Statement] -> (ReadList, WriteList) +analyzeStatementsMin = analyzeStatementsGeneric False diff --git a/src/Jikka/RestrictedPython/Language/WithoutLoc.hs b/src/Jikka/RestrictedPython/Language/WithoutLoc.hs new file mode 100644 --- /dev/null +++ b/src/Jikka/RestrictedPython/Language/WithoutLoc.hs @@ -0,0 +1,55 @@ +module Jikka.RestrictedPython.Language.WithoutLoc where + +import Jikka.Common.Location +import Jikka.RestrictedPython.Language.Expr + +constIntExp :: Integer -> Expr' +constIntExp = withoutLoc . Constant . ConstInt + +constBoolExp :: Bool -> Expr' +constBoolExp = withoutLoc . Constant . ConstBool + +constBuiltinExp :: Builtin -> Expr' +constBuiltinExp = withoutLoc . Constant . ConstBuiltin + +binOp :: Expr' -> Operator -> Expr' -> Expr' +binOp e1 op e2 = withoutLoc (BinOp e1 op e2) + +addExp :: Expr' -> Expr' -> Expr' +addExp e1 e2 = binOp e1 Add e2 + +subExp :: Expr' -> Expr' -> Expr' +subExp e1 e2 = binOp e1 Sub e2 + +multExp :: Expr' -> Expr' -> Expr' +multExp e1 e2 = binOp e1 Mult e2 + +unaryOp :: UnaryOp -> Expr' -> Expr' +unaryOp op e = withoutLoc (UnaryOp op e) + +eqExp :: Type -> Expr' -> Expr' -> Expr' +eqExp t e1 e2 = withoutLoc (Compare e1 (CmpOp' Eq' t) e2) + +name :: VarName' -> Expr' +name = withoutLoc . Name + +call :: Expr' -> [Expr'] -> Expr' +call f args = withoutLoc (Call f args) + +list :: Type -> [Expr'] -> Expr' +list = (withoutLoc .) . List + +listComp :: Expr' -> Comprehension -> Expr' +listComp = (withoutLoc .) . ListComp + +subscript :: Expr' -> Expr' -> Expr' +subscript = (withoutLoc .) . Subscript + +nameTrg :: VarName' -> Target' +nameTrg = withoutLoc . NameTrg + +subscriptTrg :: Target' -> Expr' -> Target' +subscriptTrg = (withoutLoc .) . SubscriptTrg + +tupleTrg :: [Target'] -> Target' +tupleTrg = withoutLoc . TupleTrg diff --git a/test/Jikka/CPlusPlus/Convert/FromCoreSpec.hs b/test/Jikka/CPlusPlus/Convert/FromCoreSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/CPlusPlus/Convert/FromCoreSpec.hs @@ -0,0 +1,68 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.CPlusPlus.Convert.FromCoreSpec + ( spec, + ) +where + +import Jikka.CPlusPlus.Convert.FromCore +import qualified Jikka.CPlusPlus.Language.Expr as Y +import Jikka.Common.Alpha +import Jikka.Common.Error +import qualified Jikka.Core.Language.BuiltinPatterns as X +import qualified Jikka.Core.Language.Expr as X +import Test.Hspec + +run' :: X.Program -> Either Error Y.Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + X.ToplevelLetRec + "f" + [("n", X.IntTy)] + X.IntTy + ( X.If' + X.IntTy + (X.Equal' X.IntTy (X.Var "n") X.Lit0) + X.Lit1 + (X.Mult' (X.Var "n") (X.App (X.Var "f") (X.Minus' (X.Var "n") X.Lit1))) + ) + (X.ResultExpr (X.Var "f")) + let expectedF = + Y.FunDef + Y.TyInt64 + "f_0" + [(Y.TyInt64, "n_1")] + [ Y.Declare Y.TyInt64 "x2" Y.DeclareDefault, + Y.If + (Y.BinOp Y.Equal (Y.Var "n_1") (Y.Lit (Y.LitInt64 0))) + [Y.Assign (Y.AssignExpr Y.SimpleAssign (Y.LeftVar "x2") (Y.Lit (Y.LitInt64 1)))] + ( Just + [ Y.Assign + ( Y.AssignExpr + Y.SimpleAssign + (Y.LeftVar "x2") + ( Y.BinOp + Y.Mul + (Y.Var "n_1") + ( Y.CallExpr + (Y.Var "f_0") + [Y.BinOp Y.Sub (Y.Var "n_1") (Y.Lit (Y.LitInt64 1))] + ) + ) + ) + ] + ), + Y.Return (Y.Var "x2") + ] + let expectedSolve = + Y.FunDef + Y.TyInt64 + "solve" + [(Y.TyInt64, "a3")] + [Y.Return (Y.CallExpr (Y.Var "f_0") [Y.Var "a3"])] + let expected = Y.Program [expectedF, expectedSolve] + run' prog `shouldBe` Right expected diff --git a/test/Jikka/CPlusPlus/FormatSpec.hs b/test/Jikka/CPlusPlus/FormatSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/CPlusPlus/FormatSpec.hs @@ -0,0 +1,44 @@ +module Jikka.CPlusPlus.FormatSpec + ( spec, + ) +where + +import Data.List +import Jikka.CPlusPlus.Format +import Jikka.CPlusPlus.Language.Expr +import Test.Hspec + +run'' :: Program -> [String] +run'' prog = dropWhile ("#include" `isPrefixOf`) (lines (run' prog)) + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let program = + Program + [ FunDef + TyInt64 + (VarName "solve") + [(TyInt32, VarName "n")] + [ Declare TyInt64 (VarName "x") (DeclareCopy (Lit (LitInt64 0))), + For + TyInt32 + (VarName "i") + (Lit (LitInt32 0)) + (BinOp LessThan (Var (VarName "i")) (Var (VarName "n"))) + (AssignIncr (LeftVar (VarName "i"))) + [ Assign (AssignExpr AddAssign (LeftVar (VarName "x")) (Call (Cast TyInt64) [Var (VarName "i")])) + ], + Return (Var (VarName "x")) + ] + ] + let formatted = + [ "int64_t solve(int32_t n) {", + " int64_t x = 0;", + " for (int32_t i = 0; i < n; ++ i) {", + " x += int64_t(i);", + " }", + " return x;", + "}" + ] + run'' program `shouldBe` formatted diff --git a/test/Jikka/Common/MatrixSpec.hs b/test/Jikka/Common/MatrixSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Common/MatrixSpec.hs @@ -0,0 +1,57 @@ +module Jikka.Common.MatrixSpec + ( spec, + ) +where + +import qualified Data.Vector as V +import Jikka.Common.Matrix +import Test.Hspec + +makeMatrix'' :: [[Integer]] -> Matrix Integer +makeMatrix'' = makeMatrix' . V.fromList . map V.fromList + +spec :: Spec +spec = do + describe "matcheck" $ do + it "works" $ do + let f = V.fromList $ map V.fromList [[1, 2, 3], [3, 4, 5]] + let expected = True + matcheck f `shouldBe` expected + it "works'" $ do + let f = V.fromList $ map V.fromList [[1, 2, 3], [3, 4]] + let expected = False + matcheck f `shouldBe` expected + describe "matap" $ do + it "works" $ do + let f = makeMatrix'' [[1, 2, 3], [3, 4, 5]] + let x = V.fromList [1, 2, 3] + let y = V.fromList [14, 26] + matap f x `shouldBe` y + + describe "matadd" $ do + it "works" $ do + let f = makeMatrix'' [[1, 2, 3], [3, 4, 5]] + let g = makeMatrix'' [[7, 7, 7], [6, 5, 4]] + let h = makeMatrix'' [[8, 9, 10], [9, 9, 9]] + matadd f g `shouldBe` h + + describe "matmul" $ do + it "works" $ do + let f = makeMatrix'' [[1, 2, 3], [3, 4, 5]] + let g = makeMatrix'' [[1, 2], [3, 4], [5, 6]] + let h = makeMatrix'' [[22, 28], [40, 52]] + matmul f g `shouldBe` h + + describe "matscalar" $ do + it "works" $ do + let k = 3 + let f = makeMatrix'' [[1, 2, 3], [3, 4, 5]] + let g = makeMatrix'' [[3, 6, 9], [9, 12, 15]] + matscalar k f `shouldBe` g + + describe "matpow" $ do + it "works" $ do + let f = makeMatrix'' [[1, 1], [1, 0]] + let k = 10 + let g = makeMatrix'' [[89, 55], [55, 34]] + matpow f k `shouldBe` g diff --git a/test/Jikka/Common/Parse/JoinLinesSpec.hs b/test/Jikka/Common/Parse/JoinLinesSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Common/Parse/JoinLinesSpec.hs @@ -0,0 +1,53 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Common.Parse.JoinLinesSpec + ( spec, + ) +where + +import Data.Either (isLeft) +import Jikka.Common.Error +import Jikka.Common.Location +import Jikka.Common.Parse.JoinLines +import Test.Hspec + +open :: String +open = "" + +close :: String +close = "" + +newline :: String +newline = "" + +joinLinesWithParens' :: [String] -> Either Error [String] +joinLinesWithParens' = post . go . pre + where + pre = map (WithLoc (Loc 1 0 (-1))) + go :: [WithLoc String] -> Either Error [WithLoc String] + go = joinLinesWithParens (== open) (== close) (== newline) + post = fmap (map value) + +removeEmptyLines' :: [String] -> [String] +removeEmptyLines' = post . go . pre + where + pre = map (WithLoc (Loc 0 0 (-1))) + go :: [WithLoc String] -> [WithLoc String] + go = removeEmptyLines (== newline) + post = map value + +spec :: Spec +spec = do + describe "joinLinesWithParens" $ do + it "works" $ do + let tokens = ["f", newline, open, newline, "x", ",", newline, "y", close, newline] + let expected = ["f", newline, open, "x", ",", "y", close, newline] + joinLinesWithParens' tokens `shouldBe` Right expected + it "reports unmatching parens" $ do + let tokens = ["f", open, newline, "x"] + joinLinesWithParens' tokens `shouldSatisfy` isLeft + describe "removeEmptyLines" $ do + it "works" $ do + let tokens = [newline, "x", newline, newline, "y", newline] + let expected = ["x", newline, "y", newline] + removeEmptyLines' tokens `shouldBe` expected diff --git a/test/Jikka/Common/Parse/OffsideRuleSpec.hs b/test/Jikka/Common/Parse/OffsideRuleSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Common/Parse/OffsideRuleSpec.hs @@ -0,0 +1,60 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Common.Parse.OffsideRuleSpec + ( spec, + ) +where + +import Data.Either (isLeft) +import Jikka.Common.Error (Error) +import Jikka.Common.Location +import Jikka.Common.Parse.OffsideRule (insertIndents) +import Test.Hspec + +indent :: String +indent = "" + +dedent :: String +dedent = "" + +newline :: String +newline = "" + +-- | This takes only `column` because `insertIndents` doesn't use `line` values. +token :: String -> Int -> WithLoc String +token s x = WithLoc (Loc 0 x 0) s + +run :: [WithLoc String] -> Either Error [String] +run = post . go + where + go :: [WithLoc String] -> Either Error [WithLoc String] + go = insertIndents indent dedent (== newline) + post = fmap (map value) + +spec :: Spec +spec = describe "insertIndentTokens" $ do + it "works" $ do + let tokens = + concat + [ [token "if:" 1, token newline 4], + [token "return" 5, token newline 11], + [token "else:" 1, token newline 6], + [token "return:" 5, token newline 11] + ] + let expected = + concat + [ ["if:", newline], + [indent, "return", newline], + [dedent, "else:", newline], + [indent, "return:", newline], + [dedent] + ] + run tokens `shouldBe` Right expected + it "fails on unmatching dedents" $ do + let tokens = + concat + [ [token "if:" 1, token newline 4], + [token "return" 5, token newline 11], + [token "err" 3] + ] + run tokens `shouldSatisfy` isLeft diff --git a/test/Jikka/Common/Parse/ShuntingYardSpec.hs b/test/Jikka/Common/Parse/ShuntingYardSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Common/Parse/ShuntingYardSpec.hs @@ -0,0 +1,63 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Common.Parse.ShuntingYardSpec + ( spec, + ) +where + +import Data.Either (isLeft) +import qualified Data.Map.Strict as M +import Jikka.Common.Error +import Jikka.Common.Location +import Jikka.Common.Parse.ShuntingYard (BinOpInfo (..), Fixity (..), run) +import Test.Hspec + +-- Haskell's one +builtInOps :: M.Map String BinOpInfo +builtInOps = + let op fixity prec name = (name, BinOpInfo fixity prec) + in M.fromList + [ op Rightfix 8 "**", + op Leftfix 7 "*", + op Leftfix 7 "/", + op Leftfix 7 "%", + op Leftfix 6 "+", + op Leftfix 6 "-", + op Nonfix 4 "==", + op Nonfix 4 "/=", + op Nonfix 4 "<", + op Nonfix 4 "<=", + op Nonfix 4 ">", + op Nonfix 4 ">=", + op Rightfix 3 "&&", + op Rightfix 2 "||" + ] + +run' :: [String] -> Either Error String +run' tokens = value <$> run info apply (f (map putPos tokens)) + where + info op = maybeToError (Error (show op ++ " is not defined")) $ M.lookup op builtInOps + apply op x y = putPos $ "(" ++ value x ++ " " ++ value op ++ " " ++ value y ++ ")" + f [] = error "the length of tokens must be odd" + f [z] = (z, []) + f (x : y : zs) = let (z, ws) = f zs in (x, (y, z) : ws) + putPos = WithLoc (Loc 0 0 (-1)) + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let tokens = ["a", "+", "b", "**", "c", "*", "d"] + let tree = "(a + ((b ** c) * d))" + run' tokens `shouldBe` Right tree + it "recognizes left-fixity" $ do + let tokens = ["a", "-", "b", "-", "c", "-", "d"] + let tree = "(((a - b) - c) - d)" + run' tokens `shouldBe` Right tree + it "recognizes right-fixity" $ do + let tokens = ["a", "&&", "b", "&&", "c", "&&", "d"] + let tree = "(a && (b && (c && d)))" + run' tokens `shouldBe` Right tree + it "reports the error on chained non-fix ops" $ do + let tokens = ["a", "==", "b", "==", "c"] + run' tokens `shouldSatisfy` isLeft diff --git a/test/Jikka/Core/Convert/ANormalSpec.hs b/test/Jikka/Core/Convert/ANormalSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Convert/ANormalSpec.hs @@ -0,0 +1,33 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Convert.ANormalSpec + ( spec, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Convert.ANormal (run) +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + ResultExpr $ + Plus' + (Let "x" IntTy Lit1 (Var "x")) + (App (Lam "x" IntTy (Var "x")) Lit1) + let expected = + ResultExpr $ + Let "x$0" IntTy Lit1 $ + Let "$3" (Fun1STy IntTy) (Lam "x$1" IntTy (Var "x$1")) $ + Let "$2" (Fun1STy IntTy) (Var "$3") $ + Let "$4" IntTy (App (Var "$2") Lit1) $ + Plus' (Var "x$0") (Var "$4") + run' prog `shouldBe` Right expected diff --git a/test/Jikka/Core/Convert/AlphaSpec.hs b/test/Jikka/Core/Convert/AlphaSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Convert/AlphaSpec.hs @@ -0,0 +1,47 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Convert.AlphaSpec + ( spec, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Convert.Alpha (run) +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + ResultExpr + ( Let + "x" + IntTy + Lit0 + ( Let + "x" + IntTy + (Plus' (Var "x") Lit1) + (Var "x") + ) + ) + let expected = + ResultExpr + ( Let + "x$0" + IntTy + Lit0 + ( Let + "x$1" + IntTy + (Plus' (Var "x$0") Lit1) + (Var "x$1") + ) + ) + run' prog `shouldBe` Right expected diff --git a/test/Jikka/Core/Convert/BetaSpec.hs b/test/Jikka/Core/Convert/BetaSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Convert/BetaSpec.hs @@ -0,0 +1,38 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Convert.BetaSpec + ( spec, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Convert.Beta (run) +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + ResultExpr + ( Lam + "a" + IntTy + ( App + (Lam "x" IntTy (Plus' (Var "x") (Var "x"))) + (Var "a") + ) + ) + let expected = + ResultExpr + ( Lam + "a$0" + IntTy + (Plus' (Var "a$0") (Var "a$0")) + ) + run' prog `shouldBe` Right expected diff --git a/test/Jikka/Core/Convert/CloseSumSpec.hs b/test/Jikka/Core/Convert/CloseSumSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Convert/CloseSumSpec.hs @@ -0,0 +1,51 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Convert.CloseSumSpec (spec) where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Convert.CloseSum (rule) +import qualified Jikka.Core.Convert.ConstantFolding as ConstantFolding +import qualified Jikka.Core.Convert.ShortCutFusion as ShortCutFusion +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.RewriteRules +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . applyRewriteRuleProgram' (rule <> ConstantFolding.rule <> ShortCutFusion.rule) + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + ResultExpr + ( Lam + "n" + IntTy + ( Sum' + ( Map' + IntTy + IntTy + (Lam "x" IntTy (Mult' (Lit (LitInt 100)) (Var "x"))) + (Range1' (Var "n")) + ) + ) + ) + let expected = + ResultExpr + ( Lam + "n" + IntTy + ( Mult' + (Lit (LitInt 100)) + ( FloorDiv' + ( Mult' + (Var "n") + (Minus' (Var "n") Lit1) + ) + Lit2 + ) + ) + ) + run' prog `shouldBe` Right expected diff --git a/test/Jikka/Core/Convert/ConstantFoldingSpec.hs b/test/Jikka/Core/Convert/ConstantFoldingSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Convert/ConstantFoldingSpec.hs @@ -0,0 +1,27 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Convert.ConstantFoldingSpec + ( spec, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Convert.ConstantFolding (run) +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + ResultExpr $ + Lam "x" IntTy (Plus' (Mult' (LitInt' 3) (Var "x")) (Plus' (LitInt' 2) (LitInt' 1))) + let expected = + ResultExpr $ + Lam "x" IntTy (Plus' (Mult' (LitInt' 3) (Var "x")) (LitInt' 3)) + run' prog `shouldBe` Right expected diff --git a/test/Jikka/Core/Convert/ConstantPropagationSpec.hs b/test/Jikka/Core/Convert/ConstantPropagationSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Convert/ConstantPropagationSpec.hs @@ -0,0 +1,42 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Convert.ConstantPropagationSpec + ( spec, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Convert.ConstantPropagation (run) +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + ResultExpr + ( Let + "x" + IntTy + Lit1 + ( Let + "f" + (ListTy IntTy) + (Range1' (LitInt' 100)) + (Plus' (Var "x") (Plus' (Var "x") (At' IntTy (Var "f") (Var "x")))) + ) + ) + let expected = + ResultExpr + ( Let + "f" + (ListTy IntTy) + (Range1' (LitInt' 100)) + (Plus' Lit1 (Plus' Lit1 (At' IntTy (Var "f") Lit1))) + ) + run' prog `shouldBe` Right expected diff --git a/test/Jikka/Core/Convert/EtaSpec.hs b/test/Jikka/Core/Convert/EtaSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Convert/EtaSpec.hs @@ -0,0 +1,36 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Convert.EtaSpec + ( spec, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Convert.Eta (run) +import Jikka.Core.Language.Expr +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + ResultExpr + ( Let + "plus" + (FunTy IntTy (FunTy IntTy IntTy)) + (Lit (LitBuiltin Plus)) + (Var "plus") + ) + let expected = + ResultExpr + ( Let + "plus" + (FunTy IntTy (FunTy IntTy IntTy)) + (Lam "$0" IntTy (Lam "$1" IntTy (App2 (Lit (LitBuiltin Plus)) (Var "$0") (Var "$1")))) + (Var "plus") + ) + run' prog `shouldBe` Right expected diff --git a/test/Jikka/Core/Convert/MakeScanlSpec.hs b/test/Jikka/Core/Convert/MakeScanlSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Convert/MakeScanlSpec.hs @@ -0,0 +1,83 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Convert.MakeScanlSpec + ( spec, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import qualified Jikka.Core.Convert.Beta as Beta +import Jikka.Core.Convert.MakeScanl (rule) +import qualified Jikka.Core.Convert.ShortCutFusion as ShortCutFusion +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.RewriteRules +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . applyRewriteRuleProgram' (rule <> ShortCutFusion.rule <> Beta.rule) + +spec :: Spec +spec = describe "run" $ do + it "works on a[i + 1] = f(a[i])" $ do + let prog = + ResultExpr $ + Foldl' + IntTy + (ListTy IntTy) + ( Lam2 + "a" + (ListTy IntTy) + "i" + IntTy + (SetAt' IntTy (Var "a") (Plus' (Var "i") Lit1) (Mult' Lit2 (At' IntTy (Var "a") (Var "i")))) + ) + (SetAt' IntTy (Range1' (LitInt' 100)) (LitInt' 0) (LitInt' 1)) + (Range1' (LitInt' 99)) + let expected = + ResultExpr $ + Scanl' + IntTy + IntTy + ( Lam2 + "a$0" + IntTy + "i" + IntTy + (Mult' Lit2 (Var "a$0")) + ) + (LitInt' 1) + (Range1' (LitInt' 99)) + run' prog `shouldBe` Right expected + + it "works on a[i] = f(a[i - 1])" $ do + let prog = + ResultExpr $ + Foldl' + IntTy + (ListTy IntTy) + ( Lam2 + "a" + (ListTy IntTy) + "i" + IntTy + (SetAt' IntTy (Var "a") (Var "i") (Mult' Lit2 (At' IntTy (Var "a") (Minus' (Var "i") Lit1)))) + ) + (SetAt' IntTy (Range1' (LitInt' 100)) (LitInt' 0) (LitInt' 1)) + (Range2' (LitInt' 1) (LitInt' 99)) + let expected = + ResultExpr $ + Scanl' + IntTy + IntTy + ( Lam2 + "$3" + IntTy + "$2" + IntTy + (Mult' Lit2 (Var "$3")) + ) + (LitInt' 1) + (Range1' (Minus' (LitInt' 99) (LitInt' 1))) + run' prog `shouldBe` Right expected diff --git a/test/Jikka/Core/Convert/MatrixExponentiationSpec.hs b/test/Jikka/Core/Convert/MatrixExponentiationSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Convert/MatrixExponentiationSpec.hs @@ -0,0 +1,84 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Convert.MatrixExponentiationSpec + ( spec, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Convert.MatrixExponentiation (run) +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Util +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works about matrices" $ do + let ts2 = [IntTy, IntTy] + let ts22 = [TupleTy ts2, TupleTy ts2] + let proj i = Proj' ts2 i (Var "xs") + let mkTuple ts = uncurryApp (Tuple' ts) + let letConst = Let "c" IntTy (LitInt' 10) + let k = LitInt' 100 + let base = mkTuple ts2 [LitInt' 12, LitInt' 34] + let prog = + ResultExpr + ( letConst + ( Iterate' + (TupleTy ts2) + k + ( Lam + "xs" + (TupleTy ts2) + (mkTuple ts2 [Plus' (proj 0) (Mult' (Var "c") (proj 1)), proj 0]) + ) + base + ) + ) + let expected = + ResultExpr + ( letConst + ( MatAp' + 2 + 2 + ( MatPow' + 2 + (mkTuple ts22 [mkTuple ts2 [Lit1, Var "c"], mkTuple ts2 [Lit1, Lit0]]) + k + ) + base + ) + ) + run' prog `shouldBe` Right expected + it "works about integers" $ do + let letConst = Let "c" IntTy (LitInt' 10) + let k = LitInt' 100 + let base = LitInt' 1234 + let prog = + ResultExpr + ( letConst + ( Iterate' + IntTy + k + ( Lam + "x" + IntTy + (Plus' (Mult' (Var "c") (Var "x")) (LitInt' 2)) + ) + base + ) + ) + let expected = + ResultExpr + ( letConst + ( Plus' + (Mult' (Pow' (Var "c") k) base) + (Mult' (FloorDiv' (Minus' (Pow' (Var "c") k) (LitInt' 1)) (Minus' (Var "c") (LitInt' 1))) (LitInt' 2)) + ) + ) + run' prog `shouldBe` Right expected diff --git a/test/Jikka/Core/Convert/PropagateModSpec.hs b/test/Jikka/Core/Convert/PropagateModSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Convert/PropagateModSpec.hs @@ -0,0 +1,37 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Convert.PropagateModSpec + ( spec, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Convert.PropagateMod (run) +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let m = LitInt' 1000000007 + let f e = FloorMod' e m + let prog = + ResultExpr + ( Lam + "y" + IntTy + (f (App (Lam "x" IntTy (Plus' (Mult' (Var "x") (Var "x")) (Var "x"))) (Var "y"))) + ) + let expected = + ResultExpr + ( Lam + "y" + IntTy + (App (Lam "x$0" IntTy (ModPlus' (ModMult' (f (Var "x$0")) (f (Var "x$0")) m) (f (Var "x$0")) m)) (Var "y")) + ) + run' prog `shouldBe` Right expected diff --git a/test/Jikka/Core/Convert/RemoveUnusedVarsSpec.hs b/test/Jikka/Core/Convert/RemoveUnusedVarsSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Convert/RemoveUnusedVarsSpec.hs @@ -0,0 +1,25 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Convert.RemoveUnusedVarsSpec (spec) where + +import Jikka.Core.Convert.RemoveUnusedVars (run) +import Jikka.Core.Language.Expr +import Test.Hspec + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + ToplevelLetRec + "solve" + [("x", BoolTy)] + BoolTy + (Let "y" IntTy Lit0 (Var "x")) + (ResultExpr (Var "solve")) + let expected = + ToplevelLet + "solve" + (FunTy BoolTy BoolTy) + (Lam "x" BoolTy (Var "x")) + (ResultExpr (Var "solve")) + run prog `shouldBe` Right expected diff --git a/test/Jikka/Core/Convert/ShortCutFusionSpec.hs b/test/Jikka/Core/Convert/ShortCutFusionSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Convert/ShortCutFusionSpec.hs @@ -0,0 +1,37 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Convert.ShortCutFusionSpec (spec) where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Convert.ShortCutFusion (run) +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + ResultExpr $ + Len' IntTy (Map' IntTy IntTy (Lam "n" IntTy (Mult' (LitInt' 2) (Var "n"))) (Range1' (LitInt' 100))) + let expected = + ResultExpr $ + LitInt' 100 + run' prog `shouldBe` Right expected + it "squashes foldl-map combination" $ do + let g = Lam2 "a" IntTy "i" IntTy (Plus' (Var "a") (Var "i")) + let f = Lam "j" IntTy (Plus' (Var "j") Lit1) + let prog = + ResultExpr $ + Foldl' IntTy IntTy g Lit0 $ + Map' IntTy IntTy f $ + Range1' (LitInt' 100) + let expected = + ResultExpr $ + Foldl' IntTy IntTy (Lam2 "$0" IntTy "$1" IntTy (App2 g (Var "$0") (App f (Var "$1")))) Lit0 $ + Range1' (LitInt' 100) + run' prog `shouldBe` Right expected diff --git a/test/Jikka/Core/Convert/SpecializeFoldlSpec.hs b/test/Jikka/Core/Convert/SpecializeFoldlSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Convert/SpecializeFoldlSpec.hs @@ -0,0 +1,55 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Convert.SpecializeFoldlSpec (spec) where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Convert.SpecializeFoldl (run) +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + ResultExpr + ( Lam + "n" + IntTy + ( Foldl' + IntTy + IntTy + ( Lam2 + "y" + IntTy + "x" + IntTy + (Plus' (Var "y") (Var "x")) + ) + Lit0 + (Range1' (Var "n")) + ) + ) + let expected = + ResultExpr + ( Lam + "n" + IntTy + ( Sum' + ( Cons' + IntTy + Lit0 + ( Map' + IntTy + IntTy + (Lam "x" IntTy (Var "x")) + (Range1' (Var "n")) + ) + ) + ) + ) + run' prog `shouldBe` Right expected diff --git a/test/Jikka/Core/Convert/TrivialLetEliminationSpec.hs b/test/Jikka/Core/Convert/TrivialLetEliminationSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Convert/TrivialLetEliminationSpec.hs @@ -0,0 +1,42 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Convert.TrivialLetEliminationSpec + ( spec, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Convert.TrivialLetElimination (run) +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + ResultExpr + ( Let + "f" + (FunTy IntTy IntTy) + (Lam "y" IntTy (Var "y")) + ( Let + "x" + IntTy + Lit1 + (App (Var "f") (Plus' (Var "x") (Var "x"))) + ) + ) + let expected = + ResultExpr + ( Let + "x" + IntTy + Lit1 + (App (Lam "y" IntTy (Var "y")) (Plus' (Var "x") (Var "x"))) + ) + run' prog `shouldBe` Right expected diff --git a/test/Jikka/Core/Convert/TypeInferSpec.hs b/test/Jikka/Core/Convert/TypeInferSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Convert/TypeInferSpec.hs @@ -0,0 +1,111 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Convert.TypeInferSpec + ( spec, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Convert.TypeInfer (run) +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + ResultExpr + ( Let + "x" + (VarTy "t1") + Lit0 + ( Let + "y" + (VarTy "t2") + (Plus' (Var "x") Lit1) + (Var "y") + ) + ) + let expected = + ResultExpr + ( Let + "x" + IntTy + Lit0 + ( Let + "y" + IntTy + (Plus' (Var "x") Lit1) + (Var "y") + ) + ) + run' prog `shouldBe` Right expected + it "works on let-rec" $ do + let prog = + ToplevelLetRec + "f" + [("x", VarTy "t1")] + (VarTy "t2") + (Var "x") + (ResultExpr (App (Var "f") Lit0)) + let expected = + ToplevelLetRec + "f" + [("x", IntTy)] + IntTy + (Var "x") + (ResultExpr (App (Var "f") Lit0)) + run' prog `shouldBe` Right expected + it "replaces undetermined types with 0-tuples" $ do + let prog = + ToplevelLetRec + "f" + [("x", VarTy "t1")] + (VarTy "t2") + (Var "x") + (ResultExpr Lit0) + let expected = + ToplevelLetRec + "f" + [("x", TupleTy [])] + (TupleTy []) + (Var "x") + (ResultExpr Lit0) + run' prog `shouldBe` Right expected + it "works on builtin functions" $ do + let prog = + ToplevelLetRec + "solve" + [("n", IntTy)] + IntTy + ( If' + (VarTy "$0") + (Equal' IntTy (Var "n") Lit0) + Lit1 + ( Mult' + (Var "n") + (App (Var "solve") (Minus' (Var "n") Lit1)) + ) + ) + (ResultExpr (Var "solve")) + let expected = + ToplevelLetRec + "solve" + [("n", IntTy)] + IntTy + ( If' + IntTy + (Equal' IntTy (Var "n") Lit0) + Lit1 + ( Mult' + (Var "n") + (App (Var "solve") (Minus' (Var "n") Lit1)) + ) + ) + (ResultExpr (Var "solve")) + run' prog `shouldBe` Right expected diff --git a/test/Jikka/Core/Convert/UnpackTupleSpec.hs b/test/Jikka/Core/Convert/UnpackTupleSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Convert/UnpackTupleSpec.hs @@ -0,0 +1,37 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Convert.UnpackTupleSpec + ( spec, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Convert.UnpackTuple (run) +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Util +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let ts = [IntTy, IntTy] + let prog = + ResultExpr $ + App (Lam "x" (TupleTy ts) (Plus' (Proj' ts 0 (Var "x")) (Proj' ts 1 (Var "x")))) (uncurryApp (Tuple' ts) [LitInt' 0, LitInt' 1]) + let expected = + ResultExpr $ + App2 (Lam2 "x$1" IntTy "x$2" IntTy (Plus' (Var "x$1") (Var "x$2"))) (LitInt' 0) (LitInt' 1) + run' prog `shouldBe` Right expected + it "works on foldl" $ do + let prog = + ResultExpr $ + Foldl' IntTy (TupleTy [IntTy]) (Lam2 "x" (TupleTy [IntTy]) "y" IntTy (uncurryApp (Tuple' [IntTy]) [Plus' (Proj' [IntTy] 0 (Var "x")) (Var "y")])) (uncurryApp (Tuple' [IntTy]) [LitInt' 0]) (Range1' (LitInt' 10)) + let expected = + ResultExpr $ + uncurryApp (Tuple' [IntTy]) [Foldl' IntTy IntTy (Lam2 "x$0" IntTy "y$1" IntTy (Plus' (Var "x$0") (Var "y$1"))) (LitInt' 0) (Range1' (LitInt' 10))] + run' prog `shouldBe` Right expected diff --git a/test/Jikka/Core/EvaluateSpec.hs b/test/Jikka/Core/EvaluateSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/EvaluateSpec.hs @@ -0,0 +1,68 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.EvaluateSpec (spec) where + +import qualified Data.Vector as V +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Core.Evaluate +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Jikka.Core.Language.Value (formatValue) +import Test.Hspec + +run'' :: Program -> [Value] -> Either Error Value +run'' = (flip evalAlphaT 0 .) . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + ToplevelLetRec + "solve" + [("xs", ListTy IntTy)] + IntTy + ( Plus' + (Sum' (Var "xs")) + (Len' IntTy (Var "xs")) + ) + (ResultExpr (Var "solve")) + let args = + [ ValList + ( V.fromList + [ ValInt 1, + ValInt 2, + ValInt 5 + ] + ) + ] + let expected = ValInt 11 + (formatValue <$> run'' prog args) `shouldBe` Right (formatValue expected) + it "works on a recursive function" $ do + let prog = + ToplevelLetRec + "fact" + [("n", IntTy)] + IntTy + ( App + ( If' + (FunTy UnitTy IntTy) + (Equal' IntTy (Var "n") Lit0) + (Lam "x" UnitTy Lit1) + ( Lam + "x" + UnitTy + ( Mult' + (Var "n") + (App (Var "fact") (Minus' (Var "n") Lit1)) + ) + ) + ) + (Tuple' []) + ) + (ResultExpr (Var "fact")) + let args = + [ ValInt 10 + ] + let expected = ValInt 3628800 + (formatValue <$> run'' prog args) `shouldBe` Right (formatValue expected) diff --git a/test/Jikka/Core/FormatSpec.hs b/test/Jikka/Core/FormatSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/FormatSpec.hs @@ -0,0 +1,48 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.FormatSpec + ( spec, + ) +where + +import Jikka.Core.Format +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Test.Hspec + +spec :: Spec +spec = describe "formatExpr" $ do + it "works" $ do + let program = + ToplevelLetRec + "solve$0" + [("n$1", IntTy)] + IntTy + ( Let + "xs$2" + (ListTy IntTy) + ( Map' + IntTy + IntTy + ( Lam + "i$3" + IntTy + (Mult' (Var "i$3") (Var "i$3")) + ) + (Range1' (Var "n$1")) + ) + (Sum' (Var "xs$2")) + ) + (ResultExpr (Var "solve$0")) + let expected = + unlines + [ "let rec solve$0 (n$1: int): int =", + " let xs$2: int list =", + " map((fun (i$3: int) ->", + " (i$3 * i$3)", + " ), range1(n$1))", + " in sum(xs$2)", + "in", + "solve$0" + ] + formatProgram program `shouldBe` expected diff --git a/test/Jikka/Core/Language/ArithmeticalExprSpec.hs b/test/Jikka/Core/Language/ArithmeticalExprSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Language/ArithmeticalExprSpec.hs @@ -0,0 +1,29 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Language.ArithmeticalExprSpec + ( spec, + ) +where + +import qualified Data.Vector as V +import Jikka.Core.Language.ArithmeticalExpr +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Test.Hspec + +spec :: Spec +spec = do + describe "makeVectorFromArithmeticalExpr" $ do + it "works" $ do + let xs = V.fromList ["x", "y"] + let e = + parseArithmeticalExpr + (Plus' (Var "x") (Plus' (Mult' (LitInt' 3) (Var "y")) (Minus' (Var "x") (LitInt' 10)))) + let f = V.fromList [parseArithmeticalExpr (LitInt' 2), parseArithmeticalExpr (LitInt' 3)] + let c = parseArithmeticalExpr (LitInt' (-10)) + makeVectorFromArithmeticalExpr xs e `shouldBe` Just (f, c) + describe "normalizeArithmeticalExpr" $ do + it "works" $ do + let e = Plus' (LitInt' 2) (Plus' (Var "a") (LitInt' (-2))) + let expected = Var "a" + (formatArithmeticalExpr . parseArithmeticalExpr) e `shouldBe` expected diff --git a/test/Jikka/Core/Language/BetaSpec.hs b/test/Jikka/Core/Language/BetaSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Core/Language/BetaSpec.hs @@ -0,0 +1,31 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Core.Language.BetaSpec + ( spec, + ) +where + +import Jikka.Common.Alpha +import Jikka.Core.Language.Beta +import Jikka.Core.Language.BuiltinPatterns +import Jikka.Core.Language.Expr +import Test.Hspec + +substitute' :: VarName -> Expr -> Expr -> Expr +substitute' x e = flip evalAlpha 0 . substitute x e + +spec :: Spec +spec = do + describe "substitute" $ do + it "renames scoped variables of lam if required" $ do + let x = "x" + let a = Var "y" + let e = Lam "y" IntTy (Plus' (Var "x") (Var "y")) + let expected = Lam "y$0" IntTy (Plus' (Var "y") (Var "y$0")) + substitute' x a e `shouldBe` expected + it "renames scoped variables of let if required" $ do + let x = "x" + let a = Var "y" + let e = Let "y" IntTy (Var "y") (Plus' (Var "x") (Var "y")) + let expected = Let "y$0" IntTy (Var "y") (Plus' (Var "y") (Var "y$0")) + substitute' x a e `shouldBe` expected diff --git a/test/Jikka/Python/Convert/ToRestrictedPythonSpec.hs b/test/Jikka/Python/Convert/ToRestrictedPythonSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Python/Convert/ToRestrictedPythonSpec.hs @@ -0,0 +1,66 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Python.Convert.ToRestrictedPythonSpec + ( spec, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.Common.Location +import Jikka.Python.Convert.ToRestrictedPython (run) +import qualified Jikka.Python.Language.Expr as X +import qualified Jikka.RestrictedPython.Language.Expr as Y +import Test.Hspec + +at :: a -> Int -> WithLoc a +at a x = WithLoc (Loc 0 x (-1)) a + +at' :: a -> Int -> WithLoc' a +at' a x = WithLoc' (Just (Loc 0 x (-1))) a + +run' :: X.Program -> Either Error Y.Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let parsed = + [ X.FunctionDef + ("solve" `at` 6) + (X.emptyArguments {X.argsArgs = [("x" `at` 7, Nothing)]}) + [ X.Assign [X.Name ("y" `at` 8) `at` 5] (X.Name ("x" `at` 3) `at` 3) `at` 4 + ] + [] + (Just (X.Name ("int" `at` 2) `at` 2)) + `at` 1 + ] + let expected = + [ Y.ToplevelFunctionDef + ("solve" `at'` 6) + [("x" `at'` 7, Y.VarTy "$0")] + Y.IntTy + [ Y.AnnAssign (Y.NameTrg ("y" `at'` 8) `at'` 5) (Y.VarTy "$1") (Y.Name ("x" `at'` 3) `at'` 3) + ] + ] + run' parsed `shouldBe` Right expected + it "works on recursive functions" $ do + let parsed = + [ X.FunctionDef + ("f" `at` 6) + (X.emptyArguments {X.argsArgs = [("x" `at` 7, Nothing)]}) + [ X.Return (Just (X.Call (X.Name ("f" `at` 8) `at` 5) [X.Name ("x" `at` 4) `at` 4] [] `at` 3)) `at` 2 + ] + [] + Nothing + `at` 1 + ] + let expected = + [ Y.ToplevelFunctionDef + ("f" `at'` 6) + [("x" `at'` 7, Y.VarTy "$0")] + (Y.VarTy "$1") + [ Y.Return (Y.Call (Y.Name ("f" `at'` 8) `at'` 5) [Y.Name ("x" `at'` 4) `at'` 4] `at'` 3) + ] + ] + run' parsed `shouldBe` Right expected diff --git a/test/Jikka/Python/Parse/AlexSpec.hs b/test/Jikka/Python/Parse/AlexSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Python/Parse/AlexSpec.hs @@ -0,0 +1,36 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Python.Parse.AlexSpec + ( spec, + ) +where + +import Jikka.Common.Error (Error) +import Jikka.Common.Location +import Jikka.Python.Parse.Alex (run) +import Jikka.Python.Parse.Token (Token (..)) +import Test.Hspec + +run' :: String -> Either Error [Token] +run' input = do + tokens <- run input + return $ map value tokens + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let input = "abc ** 123" + let tokens = [Ident "abc", PowOp, Int 123] + run' input `shouldBe` Right tokens + it "puts 1-based position info" $ do + let input = "abc def\n123" + let tokens = [WithLoc (Loc 1 1 3) $ Ident "abc", WithLoc (Loc 1 5 3) Def, WithLoc (Loc 1 8 1) Newline, WithLoc (Loc 2 1 3) $ Int 123] + run input `shouldBe` Right tokens + it "inserts tokens" $ do + let input = "if:\n return" + let tokens = [If, Colon, Newline, Indent, Return, Dedent] + run' input `shouldBe` Right tokens + it "uses the longest match" $ do + let input = "i in int ints" + let tokens = [Ident "i", In, Ident "int", Ident "ints"] + run' input `shouldBe` Right tokens diff --git a/test/Jikka/Python/Parse/HappySpec.hs b/test/Jikka/Python/Parse/HappySpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Python/Parse/HappySpec.hs @@ -0,0 +1,49 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Python.Parse.HappySpec + ( spec, + ) +where + +import Jikka.Common.Error (Error) +import Jikka.Common.Location +import Jikka.Python.Language.Expr +import Jikka.Python.Language.Util +import Jikka.Python.Parse.Happy +import qualified Jikka.Python.Parse.Token as L +import Test.Hspec + +at :: a -> (Int, Int) -> WithLoc a +at token (y, x) = WithLoc (Loc y x (-1)) token + +run' :: [[L.Token]] -> Either Error Program +run' tokens = run . concat $ zipWith (\y -> zipWith (\x token -> token `at` (y, x)) [1 ..]) [1 ..] tokens + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let input = + [ [L.Def, L.Ident "solve", L.OpenParen, L.CloseParen, L.Arrow, L.Ident "int", L.Colon, L.Newline], + [L.Indent, L.Return, L.Int 42, L.Newline], + [L.Dedent] + ] + let parsed = [FunctionDef ("solve" `at` (1, 2)) emptyArguments [Return (Just (constIntExp 42 `at` (2, 3))) `at` (2, 2)] [] (Just $ Name ("int" `at` (1, 6)) `at` (1, 6)) `at` (1, 1)] + run' input `shouldBe` Right parsed + it "works on a small fun def" $ do + let input = + [ [L.Def, L.Ident "solve", L.OpenParen, L.Ident "p", L.CloseParen, L.Colon, L.Newline], + [L.Indent, L.If, L.Ident "p", L.Colon, L.Newline], + [L.Indent, L.Return, L.Int 0, L.Newline], + [L.Dedent, L.Else, L.Colon, L.Newline], + [L.Indent, L.Return, L.Int 1, L.Newline], + [L.Dedent], + [L.Dedent] + ] + let parsed = + [FunctionDef ("solve" `at` (1, 2)) (emptyArguments {argsArgs = [("p" `at` (1, 4), Nothing)]}) [If (Name ("p" `at` (2, 3)) `at` (2, 3)) [Return (Just (constIntExp 0 `at` (3, 3))) `at` (3, 2)] [Return (Just (constIntExp 1 `at` (5, 3))) `at` (5, 2)] `at` (2, 2)] [] Nothing `at` (1, 1)] + run' input `shouldBe` Right parsed + it "works on a simple constant def" $ do + let input = [[L.Ident "MOD", L.Colon, L.Ident "int", L.Equal, L.Int 1000000007, L.Newline]] + let parsed = + [AnnAssign (Name ("MOD" `at` (1, 1)) `at` (1, 1)) (Name ("int" `at` (1, 3)) `at` (1, 3)) (Just (constIntExp 1000000007 `at` (1, 5))) `at` (1, 1)] + run' input `shouldBe` Right parsed diff --git a/test/Jikka/Python/ParseSpec.hs b/test/Jikka/Python/ParseSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/Python/ParseSpec.hs @@ -0,0 +1,30 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.Python.ParseSpec + ( spec, + ) +where + +import Data.Text (pack) +import Jikka.Common.Error (Error) +import Jikka.Common.Location +import Jikka.Python.Language.Expr +import Jikka.Python.Language.Util +import Jikka.Python.Parse +import Test.Hspec + +at :: a -> (Int, Int, Int) -> WithLoc a +at a (y, x, width) = WithLoc (Loc y x width) a + +run' :: [String] -> Either Error Program +run' lines = run "test.py" (pack $ unlines lines) + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let input = + [ "def solve() -> int:", + " return 42" + ] + let parsed = [FunctionDef ("solve" `at` (1, 5, 5)) emptyArguments [Return (Just (constIntExp 42 `at` (2, 12, 2))) `at` (2, 5, 6)] [] (Just (Name ("int" `at` (1, 16, 3)) `at` (1, 16, 3))) `at` (1, 1, 3)] + run' input `shouldBe` Right parsed diff --git a/test/Jikka/RestrictedPython/Convert/AlphaSpec.hs b/test/Jikka/RestrictedPython/Convert/AlphaSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/RestrictedPython/Convert/AlphaSpec.hs @@ -0,0 +1,599 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.RestrictedPython.Convert.AlphaSpec + ( spec, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.RestrictedPython.Convert.Alpha (run) +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.WithoutLoc +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let parsed = + [ ToplevelFunctionDef + "f" + [("n", IntTy)] + IntTy + [ If + (eqExp (VarTy "t") (name "n") (constIntExp 0)) + [ Return (constIntExp 1) + ] + [ Return (multExp (name "n") (call (name "f") [subExp (name "n") (constIntExp 1)])) + ] + ], + ToplevelFunctionDef + "solve" + [("n", IntTy)] + IntTy + [ Return (binOp (call (name "f") [name "n"]) FloorMod (constIntExp 1000000007)) + ] + ] + let expected = + [ ToplevelFunctionDef + "f" + [("n$0", IntTy)] + IntTy + [ If + (eqExp (VarTy "t") (name "n$0") (constIntExp 0)) + [ Return (constIntExp 1) + ] + [ Return (multExp (name "n$0") (call (name "f") [subExp (name "n$0") (constIntExp 1)])) + ] + ], + ToplevelFunctionDef + "solve" + [("n$1", IntTy)] + IntTy + [ Return (binOp (call (name "f") [name "n$1"]) FloorMod (constIntExp 1000000007)) + ] + ] + run' parsed `shouldBe` Right expected + it "fails with undefined variables" $ do + let parsed = + [ ToplevelFunctionDef + "solve" + [] + IntTy + [ Return (name "y") + ] + ] + let expected = WithWrapped "Jikka.RestrictedPython.Convert.Alpha" (WithGroup SymbolError (Error "undefined identifier: y")) + run' parsed `shouldBe` Left expected + it "doesn't rename builtin functions " $ do + let parsed = + [ ToplevelFunctionDef + "solve" + [] + IntTy + [ Return (name "range") + ] + ] + let expected = + [ ToplevelFunctionDef + "solve" + [] + IntTy + [ Return (name "range") + ] + ] + run' parsed `shouldBe` Right expected + it "distinguishes local variables in two diffrent functions" $ do + let parsed = + [ ToplevelFunctionDef + "foo" + [("x", IntTy)] + IntTy + [ AnnAssign (nameTrg "y") IntTy (name "x") + ], + ToplevelFunctionDef + "bar" + [("x", IntTy)] + IntTy + [ AnnAssign (nameTrg "y") IntTy (name "x") + ] + ] + let expected = + [ ToplevelFunctionDef + "foo" + [("x$0", IntTy)] + IntTy + [ AnnAssign (nameTrg "y$1") IntTy (name "x$0") + ], + ToplevelFunctionDef + "bar" + [("x$2", IntTy)] + IntTy + [ AnnAssign (nameTrg "y$3") IntTy (name "x$2") + ] + ] + run' parsed `shouldBe` Right expected + it "distinguishes variables in two diffrent for-loops" $ do + let parsed = + [ ToplevelFunctionDef + "solve" + [] + IntTy + [ For + (nameTrg "i") + (list IntTy []) + [ AnnAssign (nameTrg "x") IntTy (name "i") + ], + For + (nameTrg "i") + (list IntTy []) + [ AnnAssign (nameTrg "x") IntTy (name "i") + ] + ] + ] + let expected = + [ ToplevelFunctionDef + "solve" + [] + IntTy + [ For + (nameTrg "i$0") + (list IntTy []) + [ AnnAssign (nameTrg "x$1") IntTy (name "i$0") + ], + For + (nameTrg "i$2") + (list IntTy []) + [ AnnAssign (nameTrg "x$3") IntTy (name "i$2") + ] + ] + ] + run' parsed `shouldBe` Right expected + it "removes underscoes" $ do + let parsed = + [ ToplevelAnnAssign + "a" + (ListTy IntTy) + ( listComp + (constIntExp 0) + ( Comprehension + (nameTrg "_") + (call (name "range") [constIntExp 10]) + Nothing + ) + ) + ] + let expected = + [ ToplevelAnnAssign + "a" + (ListTy IntTy) + ( listComp + (constIntExp 0) + ( Comprehension + (nameTrg "$0") + (call (name "range") [constIntExp 10]) + Nothing + ) + ) + ] + run' parsed `shouldBe` Right expected + it "works on recursive functions" $ do + let parsed = + [ ToplevelFunctionDef + "f" + [("x", IntTy)] + IntTy + [ Return (call (name "f") [name "x"]) + ] + ] + let expected = + [ ToplevelFunctionDef + "f" + [("x$0", IntTy)] + IntTy + [ Return (call (name "f") [name "x$0"]) + ] + ] + run' parsed `shouldBe` Right expected + it "doesn't rename type variables" $ do + let parsed = + [ ToplevelFunctionDef + "f" + [("x", VarTy "x")] + (VarTy "f") + [ Return (call (name "f") [name "x"]) + ] + ] + let expected = + [ ToplevelFunctionDef + "f" + [("x$0", VarTy "x")] + (VarTy "f") + [ Return (call (name "f") [name "x$0"]) + ] + ] + run' parsed `shouldBe` Right expected + it "makes repeated assignments for the same variable to single-assignments for different variables" $ do + let parsed = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ AnnAssign (nameTrg "x") IntTy (constIntExp 0), + AnnAssign (nameTrg "x") IntTy (name "x"), + AnnAssign (nameTrg "x") IntTy (constIntExp 0), + AnnAssign (nameTrg "x") IntTy (name "x"), + AnnAssign (nameTrg "x") IntTy (constIntExp 0), + AnnAssign (nameTrg "x") IntTy (name "x") + ] + ] + let expected = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ AnnAssign (nameTrg "x$0") IntTy (constIntExp 0), + AnnAssign (nameTrg "x$1") IntTy (name "x$0"), + AnnAssign (nameTrg "x$2") IntTy (constIntExp 0), + AnnAssign (nameTrg "x$3") IntTy (name "x$2"), + AnnAssign (nameTrg "x$4") IntTy (constIntExp 0), + AnnAssign (nameTrg "x$5") IntTy (name "x$4") + ] + ] + run' parsed `shouldBe` Right expected + it "doesn't rename for augumented assignments" $ do + let parsed = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ AnnAssign (nameTrg "x") IntTy (constIntExp 0), + AnnAssign (nameTrg "x") IntTy (name "x"), + AugAssign (nameTrg "x") Add (constIntExp 0), + AugAssign (nameTrg "x") Add (name "x"), + AugAssign (nameTrg "x") Add (constIntExp 0), + AugAssign (nameTrg "x") Add (name "x") + ] + ] + let expected = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ AnnAssign (nameTrg "x$0") IntTy (constIntExp 0), + AnnAssign (nameTrg "x$1") IntTy (name "x$0"), + AugAssign (nameTrg "x$1") Add (constIntExp 0), + AugAssign (nameTrg "x$1") Add (name "x$1"), + AugAssign (nameTrg "x$1") Add (constIntExp 0), + AugAssign (nameTrg "x$1") Add (name "x$1") + ] + ] + run' parsed `shouldBe` Right expected + it "blames leaks of loop counters of for-statements" $ do + let parsed = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ For + (nameTrg "i") + (list IntTy []) + [ Return (name "i") + ], + Return (name "i") + ] + ] + let expected = WithWrapped "Jikka.RestrictedPython.Convert.Alpha" (WithGroup SymbolError (Error "undefined identifier: i")) + run' parsed `shouldBe` Left expected + it "blames leaks of loop counters of for-statements even if variables with the same names are defined" $ do + let parsed = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ AnnAssign (nameTrg "i") IntTy (constIntExp 0), + For + (nameTrg "i") + (list IntTy []) + [ Return (name "i") + ], + Return (name "i") + ] + ] + let expected = WithWrapped "Jikka.RestrictedPython.Convert.Alpha" (WithGroup SemanticError (Error "cannot redefine variable: i")) + run' parsed `shouldBe` Left expected + it "blames undefined variables which will be defined in the rest of the same loop" $ do + let parsed = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ For + (nameTrg "i") + (list IntTy []) + [ Return (name "a"), + AnnAssign (nameTrg "a") IntTy (constIntExp 0) + ] + ] + ] + let expected = WithWrapped "Jikka.RestrictedPython.Convert.Alpha" (WithGroup SymbolError (Error "undefined identifier: a")) + run' parsed `shouldBe` Left expected + it "blames using variables which are defined in only one branch of if-statement" $ do + let parsed = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ If + (constBoolExp True) + [ AnnAssign (nameTrg "a") IntTy (constIntExp 0) + ] + [], + Return (name "a") + ] + ] + let expected = WithWrapped "Jikka.RestrictedPython.Convert.Alpha" (WithGroup SymbolError (Error "undefined identifier: a")) + run' parsed `shouldBe` Left expected + it "works with variables which are defined in both branches of if-statement" $ do + let parsed = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ If + (constBoolExp True) + [ AnnAssign (nameTrg "a") IntTy (constIntExp 0) + ] + [ AnnAssign (nameTrg "a") IntTy (constIntExp 1) + ], + Return (name "a") + ] + ] + let expected = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ If + (constBoolExp True) + [ AnnAssign (nameTrg "a$0") IntTy (constIntExp 0) + ] + [ AnnAssign (nameTrg "a$0") IntTy (constIntExp 1) + ], + Return (name "a$0") + ] + ] + run' parsed `shouldBe` Right expected + it "doesn't leak loop counters of for-exprs" $ do + let parsed = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ AnnAssign (nameTrg "i") IntTy (constIntExp 0), + AnnAssign (nameTrg "a") (ListTy IntTy) (listComp (constIntExp 0) (Comprehension (nameTrg "i") (list IntTy []) Nothing)), + Return (name "i") + ] + ] + let expected = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ AnnAssign (nameTrg "i$0") IntTy (constIntExp 0), + AnnAssign (nameTrg "a$2") (ListTy IntTy) (listComp (constIntExp 0) (Comprehension (nameTrg "i$1") (list IntTy []) Nothing)), + Return (name "i$0") + ] + ] + run' parsed `shouldBe` Right expected + it "works with if-statements without else-clause" $ do + let parsed = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ AnnAssign (nameTrg "x") IntTy (constIntExp 0), + AnnAssign (nameTrg "x") IntTy (name "x"), + AnnAssign (nameTrg "x") IntTy (constIntExp 0), + AnnAssign (nameTrg "x") IntTy (name "x"), + If + (name "x") + [ AnnAssign (nameTrg "x") IntTy (constIntExp 0), + AnnAssign (nameTrg "x") IntTy (name "x"), + AnnAssign (nameTrg "y") IntTy (constIntExp 0), + AnnAssign (nameTrg "y") IntTy (name "y"), + AnnAssign (nameTrg "x") IntTy (constIntExp 0), + AnnAssign (nameTrg "x") IntTy (name "x") + ] + [], + AnnAssign (nameTrg "x") IntTy (name "x"), + AnnAssign (nameTrg "x") IntTy (constIntExp 0), + AnnAssign (nameTrg "x") IntTy (name "x") + ] + ] + let expected = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ AnnAssign (nameTrg "x$0") IntTy (constIntExp 0), + AnnAssign (nameTrg "x$1") IntTy (name "x$0"), + AnnAssign (nameTrg "x$2") IntTy (constIntExp 0), + AnnAssign (nameTrg "x$3") IntTy (name "x$2"), + If + (name "x$3") + [ AnnAssign (nameTrg "x$3") IntTy (constIntExp 0), + AnnAssign (nameTrg "x$3") IntTy (name "x$3"), + AnnAssign (nameTrg "y$4") IntTy (constIntExp 0), + AnnAssign (nameTrg "y$5") IntTy (name "y$4"), + AnnAssign (nameTrg "x$3") IntTy (constIntExp 0), + AnnAssign (nameTrg "x$3") IntTy (name "x$3") + ] + [], + AnnAssign (nameTrg "x$6") IntTy (name "x$3"), + AnnAssign (nameTrg "x$7") IntTy (constIntExp 0), + AnnAssign (nameTrg "x$8") IntTy (name "x$7") + ] + ] + run' parsed `shouldBe` Right expected + it "works with if-statements with else-clause" $ do + let parsed = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ AnnAssign (nameTrg "x") IntTy (constIntExp 0), + AnnAssign (nameTrg "x") IntTy (name "x"), + If + (name "x") + [ AugAssign (nameTrg "x") Add (name "x") + ] + [ AnnAssign (nameTrg "x") IntTy (name "x"), + AnnAssign (nameTrg "x") IntTy (name "x") + ], + AnnAssign (nameTrg "x") IntTy (name "x"), + AnnAssign (nameTrg "x") IntTy (constIntExp 0), + AnnAssign (nameTrg "x") IntTy (name "x") + ] + ] + let expected = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ AnnAssign (nameTrg "x$0") IntTy (constIntExp 0), + AnnAssign (nameTrg "x$1") IntTy (name "x$0"), + If + (name "x$1") + [ AugAssign (nameTrg "x$1") Add (name "x$1") + ] + [ AnnAssign (nameTrg "x$1") IntTy (name "x$1"), + AnnAssign (nameTrg "x$1") IntTy (name "x$1") + ], + AnnAssign (nameTrg "x$2") IntTy (name "x$1"), + AnnAssign (nameTrg "x$3") IntTy (constIntExp 0), + AnnAssign (nameTrg "x$4") IntTy (name "x$3") + ] + ] + run' parsed `shouldBe` Right expected + it "isn't confused by conflicts between variables in if-statements and variables in toplevel" $ do + let parsed = + [ ToplevelFunctionDef + "x" + [] + IntTy + [ Return (constIntExp 0) + ], + ToplevelFunctionDef + "main" + [] + IntTy + [ If + (constBoolExp True) + [ AnnAssign (nameTrg "x") IntTy (constIntExp 0) + ] + [ AnnAssign (nameTrg "x") IntTy (constIntExp 1) + ], + AnnAssign (nameTrg "x") IntTy (name "x") + ] + ] + let expected = + [ ToplevelFunctionDef + "x" + [] + IntTy + [ Return (constIntExp 0) + ], + ToplevelFunctionDef + "main" + [] + IntTy + [ If + (constBoolExp True) + [ AnnAssign (nameTrg "x$0") IntTy (constIntExp 0) + ] + [ AnnAssign (nameTrg "x$0") IntTy (constIntExp 1) + ], + AnnAssign (nameTrg "x$1") IntTy (name "x$0") + ] + ] + run' parsed `shouldBe` Right expected + it "works with for-loops" $ do + let parsed = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ AnnAssign (nameTrg "x") IntTy (constIntExp 0), + AnnAssign (nameTrg "x") IntTy (name "x"), + AnnAssign (nameTrg "x") IntTy (constIntExp 0), + AnnAssign (nameTrg "x") IntTy (name "x"), + For + (nameTrg "i") + (list IntTy []) + [ AnnAssign (nameTrg "x") IntTy (constIntExp 0), + AnnAssign (nameTrg "x") IntTy (name "x"), + AnnAssign (nameTrg "y") IntTy (constIntExp 0), + AnnAssign (nameTrg "y") IntTy (name "y"), + AnnAssign (nameTrg "x") IntTy (constIntExp 0), + AnnAssign (nameTrg "x") IntTy (name "x") + ], + AnnAssign (nameTrg "x") IntTy (name "x"), + AnnAssign (nameTrg "x") IntTy (constIntExp 0), + AnnAssign (nameTrg "x") IntTy (name "x") + ] + ] + let expected = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ AnnAssign (nameTrg "x$0") IntTy (constIntExp 0), + AnnAssign (nameTrg "x$1") IntTy (name "x$0"), + AnnAssign (nameTrg "x$2") IntTy (constIntExp 0), + AnnAssign (nameTrg "x$3") IntTy (name "x$2"), + For + (nameTrg "i$4") + (list IntTy []) + [ AnnAssign (nameTrg "x$3") IntTy (constIntExp 0), + AnnAssign (nameTrg "x$3") IntTy (name "x$3"), + AnnAssign (nameTrg "y$5") IntTy (constIntExp 0), + AnnAssign (nameTrg "y$6") IntTy (name "y$5"), + AnnAssign (nameTrg "x$3") IntTy (constIntExp 0), + AnnAssign (nameTrg "x$3") IntTy (name "x$3") + ], + AnnAssign (nameTrg "x$7") IntTy (name "x$3"), + AnnAssign (nameTrg "x$8") IntTy (constIntExp 0), + AnnAssign (nameTrg "x$9") IntTy (name "x$8") + ] + ] + run' parsed `shouldBe` Right expected + it "doesn't rename subscripted assignments" $ do + let parsed = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ AnnAssign (nameTrg "a") (ListTy IntTy) (list IntTy []), + AnnAssign (subscriptTrg (nameTrg "a") (constIntExp 0)) IntTy (subscript (name "a") (constIntExp 0)), + AnnAssign (subscriptTrg (nameTrg "a") (constIntExp 0)) IntTy (subscript (name "a") (constIntExp 0)), + AnnAssign (subscriptTrg (nameTrg "a") (constIntExp 0)) IntTy (subscript (name "a") (constIntExp 0)), + AnnAssign (subscriptTrg (nameTrg "a") (constIntExp 0)) IntTy (subscript (name "a") (constIntExp 0)) + ] + ] + let expected = + [ ToplevelFunctionDef + "main" + [] + IntTy + [ AnnAssign (nameTrg "a$0") (ListTy IntTy) (list IntTy []), + AnnAssign (subscriptTrg (nameTrg "a$0") (constIntExp 0)) IntTy (subscript (name "a$0") (constIntExp 0)), + AnnAssign (subscriptTrg (nameTrg "a$0") (constIntExp 0)) IntTy (subscript (name "a$0") (constIntExp 0)), + AnnAssign (subscriptTrg (nameTrg "a$0") (constIntExp 0)) IntTy (subscript (name "a$0") (constIntExp 0)), + AnnAssign (subscriptTrg (nameTrg "a$0") (constIntExp 0)) IntTy (subscript (name "a$0") (constIntExp 0)) + ] + ] + run' parsed `shouldBe` Right expected diff --git a/test/Jikka/RestrictedPython/Convert/RemoveUnbalancedIfSpec.hs b/test/Jikka/RestrictedPython/Convert/RemoveUnbalancedIfSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/RestrictedPython/Convert/RemoveUnbalancedIfSpec.hs @@ -0,0 +1,39 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.RestrictedPython.Convert.RemoveUnbalancedIfSpec + ( spec, + ) +where + +import Jikka.RestrictedPython.Convert.RemoveUnbalancedIf (run) +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Util +import Jikka.RestrictedPython.Language.WithoutLoc +import Test.Hspec + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + toplevelMainDef + [ If + (constBoolExp True) + [ Return (constIntExp 0) + ] + [ AnnAssign (nameTrg "a") IntTy (constIntExp 0) + ], + AnnAssign (nameTrg "b") IntTy (constIntExp 1), + Return (constIntExp 2) + ] + let expected = + toplevelMainDef + [ If + (constBoolExp True) + [ Return (constIntExp 0) + ] + [ AnnAssign (nameTrg "a") IntTy (constIntExp 0), + AnnAssign (nameTrg "b") IntTy (constIntExp 1), + Return (constIntExp 2) + ] + ] + run prog `shouldBe` expected diff --git a/test/Jikka/RestrictedPython/Convert/RemoveUnreachableSpec.hs b/test/Jikka/RestrictedPython/Convert/RemoveUnreachableSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/RestrictedPython/Convert/RemoveUnreachableSpec.hs @@ -0,0 +1,41 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.RestrictedPython.Convert.RemoveUnreachableSpec + ( spec, + ) +where + +import Jikka.RestrictedPython.Convert.RemoveUnreachable (run) +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Util +import Jikka.RestrictedPython.Language.WithoutLoc +import Test.Hspec + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + toplevelMainDef + [ AnnAssign (nameTrg "a") IntTy (constIntExp 0), + If + (constBoolExp True) + [ AnnAssign (nameTrg "b") IntTy (constIntExp 0), + Return (name "a"), + AugAssign (nameTrg "b") Add (name "1") + ] + [ Return (constIntExp 1) + ], + AugAssign (nameTrg "a") Add (constIntExp 1) + ] + let expected = + toplevelMainDef + [ AnnAssign (nameTrg "a") IntTy (constIntExp 0), + If + (constBoolExp True) + [ AnnAssign (nameTrg "b") IntTy (constIntExp 0), + Return (name "a") + ] + [ Return (constIntExp 1) + ] + ] + run prog `shouldBe` expected diff --git a/test/Jikka/RestrictedPython/Convert/ResolveBuiltinSpec.hs b/test/Jikka/RestrictedPython/Convert/ResolveBuiltinSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/RestrictedPython/Convert/ResolveBuiltinSpec.hs @@ -0,0 +1,34 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.RestrictedPython.Convert.ResolveBuiltinSpec + ( spec, + ) +where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.RestrictedPython.Convert.ResolveBuiltin (run) +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Util +import Jikka.RestrictedPython.Language.WithoutLoc +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + toplevelMainDef + [ Return (call (name "max") [list IntTy [constIntExp 2, constIntExp 3]]), + Return (call (name "max") [constIntExp 2, constIntExp 3]), + Return (call (name "max") [constIntExp 2, constIntExp 3, constIntExp 4]) + ] + let expected = + toplevelMainDef + [ Return (call (constBuiltinExp (BuiltinMax1 (VarTy "$0"))) [list IntTy [constIntExp 2, constIntExp 3]]), + Return (call (constBuiltinExp (BuiltinMax (VarTy "$1") 2)) [constIntExp 2, constIntExp 3]), + Return (call (constBuiltinExp (BuiltinMax (VarTy "$2") 3)) [constIntExp 2, constIntExp 3, constIntExp 4]) + ] + run' prog `shouldBe` Right expected diff --git a/test/Jikka/RestrictedPython/Convert/SplitLoopsSpec.hs b/test/Jikka/RestrictedPython/Convert/SplitLoopsSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/RestrictedPython/Convert/SplitLoopsSpec.hs @@ -0,0 +1,45 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.RestrictedPython.Convert.SplitLoopsSpec + ( spec, + ) +where + +import Jikka.RestrictedPython.Convert.SplitLoops (run') +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Util (toplevelMainDef) +import Jikka.RestrictedPython.Language.WithoutLoc +import Test.Hspec + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + toplevelMainDef + [ AnnAssign (nameTrg "a") IntTy (constIntExp 0), + AnnAssign (nameTrg "b") IntTy (constIntExp 0), + For + (nameTrg "i") + (call (name "range") [constIntExp 10]) + [ AnnAssign (nameTrg "c") IntTy (name "b"), + AugAssign (nameTrg "a") Add (name "i"), + AugAssign (nameTrg "b") Add (name "c") + ] + ] + let expected = + toplevelMainDef + [ AnnAssign (nameTrg "a") IntTy (constIntExp 0), + AnnAssign (nameTrg "b") IntTy (constIntExp 0), + For + (nameTrg "i") + (call (name "range") [constIntExp 10]) + [ AnnAssign (nameTrg "c") IntTy (name "b"), + AugAssign (nameTrg "b") Add (name "c") + ], + For + (nameTrg "i") + (call (name "range") [constIntExp 10]) + [ AugAssign (nameTrg "a") Add (name "i") + ] + ] + run' prog `shouldBe` expected diff --git a/test/Jikka/RestrictedPython/Convert/ToCoreSpec.hs b/test/Jikka/RestrictedPython/Convert/ToCoreSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/RestrictedPython/Convert/ToCoreSpec.hs @@ -0,0 +1,130 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.RestrictedPython.Convert.ToCoreSpec (spec) where + +import Jikka.Common.Alpha +import Jikka.Common.Error +import qualified Jikka.Core.Format as Y (formatProgram) +import qualified Jikka.Core.Language.BuiltinPatterns as Y +import qualified Jikka.Core.Language.Expr as Y +import Jikka.RestrictedPython.Convert.ToCore (run) +import qualified Jikka.RestrictedPython.Language.Expr as X +import qualified Jikka.RestrictedPython.Language.WithoutLoc as X +import Test.Hspec + +run' :: X.Program -> Either Error Y.Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let prog = + [ X.ToplevelFunctionDef + "solve" + [("n", X.IntTy)] + X.IntTy + [ X.If + (X.eqExp X.IntTy (X.name "n") (X.constIntExp 0)) + [ X.Return (X.constIntExp 1) + ] + [ X.Return (X.binOp (X.name "n") X.Mult (X.call (X.name "solve") [X.binOp (X.name "n") X.Sub (X.constIntExp 1)])) + ] + ] + ] + let expected = + Y.ToplevelLetRec + "solve" + [("n", Y.IntTy)] + Y.IntTy + ( Y.If' + (Y.VarTy "$0") + (Y.Equal' Y.IntTy (Y.Var "n") Y.Lit0) + Y.Lit1 + ( Y.Mult' + (Y.Var "n") + (Y.App (Y.Var "solve") (Y.Minus' (Y.Var "n") Y.Lit1)) + ) + ) + (Y.ResultExpr (Y.Var "solve")) + run' prog `shouldBe` Right expected + it "converts for-loops to foldl" $ do + let prog = + [ X.ToplevelFunctionDef + "solve" + [("n", X.IntTy)] + X.IntTy + [ X.AnnAssign (X.nameTrg "a") X.IntTy (X.constIntExp 0), + X.AnnAssign (X.nameTrg "b") X.IntTy (X.constIntExp 1), + X.For + (X.nameTrg "i") + (X.call (X.constBuiltinExp X.BuiltinRange1) [X.name "n"]) + [ X.AnnAssign (X.nameTrg "c") X.IntTy (X.binOp (X.name "a") X.Add (X.name "b")), + X.AnnAssign (X.nameTrg "a") X.IntTy (X.name "b"), + X.AnnAssign (X.nameTrg "b") X.IntTy (X.name "c") + ], + X.Return (X.name "a") + ] + ] + let expected = + unlines + [ "let rec solve (n: int): int =", + " let a: $0 =", + " 0", + " in let b: $1 =", + " 1", + " in let $4: ($5 * $6) =", + " foldl((fun ($4: ($5 * $6)) ($3: $2) ->", + " let b: $5 =", + " proj0($4)", + " in let a: $6 =", + " proj1($4)", + " in let i: $7 =", + " $3", + " in let c: $8 =", + " (a + b)", + " in let a: $9 =", + " b", + " in let b: $10 =", + " c", + " in tuple(b, a)", + " ), tuple(b, a), range1(n))", + " in let b: $5 =", + " proj0($4)", + " in let a: $6 =", + " proj1($4)", + " in a", + "in", + "solve" + ] + (Y.formatProgram <$> run' prog) `shouldBe` Right expected + it "converts if-statements correctly" $ do + let prog = + [ X.ToplevelFunctionDef + "solve" + [] + X.IntTy + [ X.If + (X.constBoolExp True) + [ X.AnnAssign (X.nameTrg "x") X.IntTy (X.constIntExp 1) + ] + [ X.AnnAssign (X.nameTrg "x") X.IntTy (X.constIntExp 0) + ], + X.Return (X.name "x") + ] + ] + let expected = + unlines + [ "let rec solve : int =", + " let $2: ($1,) =", + " (if true then let x: $3 =", + " 1", + " in tuple(x) else let x: $5 =", + " 0", + " in tuple(x))", + " in let x: $1 =", + " proj0($2)", + " in x", + "in", + "solve" + ] + (Y.formatProgram <$> run' prog) `shouldBe` Right expected diff --git a/test/Jikka/RestrictedPython/Convert/TypeInferSpec.hs b/test/Jikka/RestrictedPython/Convert/TypeInferSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/RestrictedPython/Convert/TypeInferSpec.hs @@ -0,0 +1,76 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.RestrictedPython.Convert.TypeInferSpec (spec) where + +import qualified Data.Map.Strict as M +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.RestrictedPython.Convert.TypeInfer +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.WithoutLoc +import Test.Hspec + +run' :: Program -> Either Error Program +run' = flip evalAlphaT 0 . run + +spec :: Spec +spec = do + describe "subst" $ do + it "works" $ do + let sigma = Subst (M.fromList [("t1", VarTy "t2"), ("t2", IntTy)]) + subst sigma (VarTy "t1") `shouldBe` IntTy + describe "run" $ do + it "works" $ do + let parsed = + [ ToplevelFunctionDef + "solve" + [("x", VarTy "t1")] + (VarTy "t2") + [ Return (binOp (name "x") Add (name "x")) + ] + ] + let expected = + [ ToplevelFunctionDef + "solve" + [("x", IntTy)] + IntTy + [ Return (binOp (name "x") Add (name "x")) + ] + ] + run' parsed `shouldBe` Right expected + it "makes undetermined type variables to the unit type" $ do + let parsed = + [ ToplevelFunctionDef + "solve" + [("x", VarTy "t1")] + (VarTy "t2") + [ Return (name "x") + ] + ] + let expected = + [ ToplevelFunctionDef + "solve" + [("x", TupleTy [])] + (TupleTy []) + [ Return (name "x") + ] + ] + run' parsed `shouldBe` Right expected + it "works on builtin functions" $ do + let parsed = + [ ToplevelFunctionDef + "solve" + [] + (ListTy BoolTy) + [ Return (withoutLoc (Call (constBuiltinExp (BuiltinMap [VarTy "t1"] (VarTy "t2"))) [withoutLoc (Lambda [("x", VarTy "t3")] (name "x")), withoutLoc (List (VarTy "t4") [])])) + ] + ] + let expected = + [ ToplevelFunctionDef + "solve" + [] + (ListTy BoolTy) + [ Return (withoutLoc (Call (constBuiltinExp (BuiltinMap [BoolTy] BoolTy)) [withoutLoc (Lambda [("x", BoolTy)] (name "x")), withoutLoc (List BoolTy [])])) + ] + ] + run' parsed `shouldBe` Right expected diff --git a/test/Jikka/RestrictedPython/EvaluateSpec.hs b/test/Jikka/RestrictedPython/EvaluateSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/RestrictedPython/EvaluateSpec.hs @@ -0,0 +1,45 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.RestrictedPython.EvaluateSpec (spec) where + +import Jikka.RestrictedPython.Evaluate +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Value +import Jikka.RestrictedPython.Language.WithoutLoc +import Test.Hspec + +spec :: Spec +spec = describe "run" $ do + it "works with recursion" $ do + let prog = + [ ToplevelFunctionDef + "solve" + [("n", IntTy)] + IntTy + [ If + (eqExp IntTy (name "n") (constIntExp 0)) + [Return (constIntExp 1)] + [Return (binOp (name "n") Mult (call (name "solve") [binOp (name "n") Sub (constIntExp 1)]))] + ] + ] + let args = [IntVal 10] + let expected = IntVal 3628800 + run prog args `shouldBe` Right expected + it "works with for-loop and assignment" $ do + let prog = + [ ToplevelFunctionDef + "solve" + [("n", IntTy)] + IntTy + [ AnnAssign (nameTrg "a") (ListTy IntTy) (call (constBuiltinExp (BuiltinList IntTy)) [call (constBuiltinExp BuiltinRange1) [name "n"]]), + For + (tupleTrg [nameTrg "i", nameTrg "a_i"]) + (call (constBuiltinExp (BuiltinEnumerate IntTy)) [name "a"]) + [ AugAssign (subscriptTrg (nameTrg "a") (name "i")) Mult (name "a_i") + ], + Return (call (constBuiltinExp BuiltinSum) [name "a"]) + ] + ] + let args = [IntVal 100] + let expected = IntVal 328350 + run prog args `shouldBe` Right expected diff --git a/test/Jikka/RestrictedPython/FormatSpec.hs b/test/Jikka/RestrictedPython/FormatSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/RestrictedPython/FormatSpec.hs @@ -0,0 +1,29 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.RestrictedPython.FormatSpec + ( spec, + ) +where + +import Jikka.RestrictedPython.Format +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.WithoutLoc +import Test.Hspec + +spec :: Spec +spec = describe "run" $ do + it "works" $ do + let program = + [ ToplevelFunctionDef + "solve$0" + [("x$1", IntTy)] + (VarTy "t$1") + [ Return (unaryOp USub (name "x$1")) + ] + ] + let formatted = + unlines + [ "def solve$0(x$1: int) -> t$1:", + " return - x$1" + ] + run' program `shouldBe` formatted diff --git a/test/Jikka/RestrictedPython/Language/BuiltinSpec.hs b/test/Jikka/RestrictedPython/Language/BuiltinSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/RestrictedPython/Language/BuiltinSpec.hs @@ -0,0 +1,62 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.RestrictedPython.Language.BuiltinSpec + ( spec, + ) +where + +import qualified Data.Set as S +import Jikka.Common.Alpha +import Jikka.Common.Error +import Jikka.RestrictedPython.Language.Builtin +import Jikka.RestrictedPython.Language.Expr +import Test.Hspec + +resolveBuiltin' :: VarName' -> Int -> Either Error Expr' +resolveBuiltin' x n = flip evalAlphaT 0 $ resolveBuiltin x n + +resolveAttribute'' :: Attribute' -> Either Error Attribute' +resolveAttribute'' x = flip evalAlphaT 0 $ resolveAttribute' x + +spec :: Spec +spec = do + describe "resolveBuiltin" $ do + it "works" $ do + let f = Right . withoutLoc . Constant . ConstBuiltin + resolveBuiltin' (withoutLoc "max") 1 `shouldBe` f (BuiltinMax1 (VarTy "$0")) + resolveBuiltin' (withoutLoc "max") 2 `shouldBe` f (BuiltinMax (VarTy "$0") 2) + resolveBuiltin' (withoutLoc "mox") 2 `shouldBe` Right (withoutLoc (Name "mox")) + it "is exhaustive" $ do + let resolve x = + let f n = resolveBuiltin' (withoutLoc x) n + in map f [0 .. 4] + let isBuiltin = \case + Left _ -> False + Right x -> case value' x of + Constant (ConstBuiltin _) -> True + _ -> False + resolve "foo" `shouldNotSatisfy` any isBuiltin + resolve "bar" `shouldNotSatisfy` any isBuiltin + resolve "sum" `shouldSatisfy` any isBuiltin + resolve "max" `shouldSatisfy` any isBuiltin + forM_ (S.toList builtinNames) $ \x -> do + resolve x `shouldSatisfy` any isBuiltin + + describe "resolveAttribute'" $ do + let resolve = resolveAttribute'' . withoutLoc . UnresolvedAttribute + it "works" $ do + resolve "count" `shouldBe` Right (withoutLoc (BuiltinCount (VarTy "$0"))) + resolve "index" `shouldBe` Right (withoutLoc (BuiltinIndex (VarTy "$0"))) + it "is exhaustive" $ do + let isBuiltin = \case + Left _ -> False + Right x -> case value' x of + UnresolvedAttribute _ -> False + _ -> True + resolve "foo" `shouldNotSatisfy` isBuiltin + resolve "bar" `shouldNotSatisfy` isBuiltin + resolve "count" `shouldSatisfy` isBuiltin + resolve "index" `shouldSatisfy` isBuiltin + forM_ (S.toList attributeNames) $ \x -> do + resolve x `shouldSatisfy` isBuiltin diff --git a/test/Jikka/RestrictedPython/Language/LintSpec.hs b/test/Jikka/RestrictedPython/Language/LintSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/RestrictedPython/Language/LintSpec.hs @@ -0,0 +1,88 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.RestrictedPython.Language.LintSpec (spec) where + +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Lint +import Jikka.RestrictedPython.Language.Util (toplevelMainDef) +import Jikka.RestrictedPython.Language.WithoutLoc +import Test.Hspec + +spec :: Spec +spec = do + describe "hasSubscriptionInLoopCounters" $ do + it "works on for-statements" $ do + let prog = + toplevelMainDef + [ For + (subscriptTrg (nameTrg "a") (constIntExp 0)) + (call (name "range") [constIntExp 100]) + [] + ] + let expected = True + hasSubscriptionInLoopCounters prog `shouldBe` expected + it "works on for-exprs" $ do + let prog = + toplevelMainDef + [ Return (listComp (constIntExp 0) (Comprehension (subscriptTrg (nameTrg "a") (constIntExp 0)) (call (name "range") [constIntExp 100]) Nothing)) + ] + let expected = True + hasSubscriptionInLoopCounters prog `shouldBe` expected + describe "hasAssignmentToLoopCounters" $ do + it "works" $ do + let prog = + toplevelMainDef + [ For + (nameTrg "i") + (call (name "range") [constIntExp 100]) + [ AugAssign (nameTrg "i") Add (constIntExp 1) + ] + ] + let expected = True + hasAssignmentToLoopCounters prog `shouldBe` expected + describe "hasAssignmentToLoopIterators" $ do + it "works" $ do + let prog = + toplevelMainDef + [ AnnAssign (nameTrg "a") (ListTy IntTy) (call (name "range") [constIntExp 100]), + For + (nameTrg "i") + (name "a") + [ AnnAssign (subscriptTrg (nameTrg "a") (constIntExp 5)) IntTy (name "i") + ] + ] + let expected = True + hasAssignmentToLoopIterators prog `shouldBe` expected + it "works even if side effects are not trivial" $ do + let prog = + toplevelMainDef + [ AnnAssign (nameTrg "a") IntTy (constIntExp 0), + For + (nameTrg "i") + (call (name "f") [name "a"]) + [ AugAssign (nameTrg "a") Add (name "i") + ] + ] + let expected = True + hasAssignmentToLoopIterators prog `shouldBe` expected + describe "hasReturnInLoops" $ do + it "works" $ do + let prog = + toplevelMainDef + [ AnnAssign (nameTrg "a") (ListTy IntTy) (call (name "range") [constIntExp 10]), + For + (nameTrg "i") + (name "a") + [ Return (constIntExp 0) + ] + ] + let expected = True + hasReturnInLoops prog `shouldBe` expected + describe "hasMixedAssignment" $ do + it "works" $ do + let prog = + toplevelMainDef + [ AnnAssign (tupleTrg [nameTrg "a", subscriptTrg (nameTrg "b") (constIntExp 0)]) (ListTy IntTy) (call (name "range") [constIntExp 10]) + ] + let expected = True + hasMixedAssignment prog `shouldBe` expected diff --git a/test/Jikka/RestrictedPython/Language/TypeInferSpec.hs b/test/Jikka/RestrictedPython/Language/TypeInferSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/RestrictedPython/Language/TypeInferSpec.hs @@ -0,0 +1,9 @@ +module Jikka.RestrictedPython.Language.TypeInferSpec + ( spec, + ) +where + +import Test.Hspec + +spec :: Spec +spec = return () diff --git a/test/Jikka/RestrictedPython/Language/UtilSpec.hs b/test/Jikka/RestrictedPython/Language/UtilSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/RestrictedPython/Language/UtilSpec.hs @@ -0,0 +1,38 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.RestrictedPython.Language.UtilSpec + ( spec, + ) +where + +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.Util +import Jikka.RestrictedPython.Language.WithoutLoc +import Test.Hspec + +spec :: Spec +spec = describe "doesAlwaysReturn" $ do + it "works" $ do + let stmt = AnnAssign (nameTrg "a") IntTy (constIntExp 0) + let expected = False + doesAlwaysReturn stmt `shouldBe` expected + it "works'" $ do + let stmt = Return (name "a") + let expected = True + doesAlwaysReturn stmt `shouldBe` expected + it "returns true for an if-statement which both branches always return" $ do + let stmt = + If + (constBoolExp True) + [ AnnAssign (nameTrg "b") IntTy (constIntExp 0), + Return (name "a"), + AugAssign (nameTrg "b") Add (name "1") + ] + [ Return (constIntExp 1) + ] + let expected = True + doesAlwaysReturn stmt `shouldBe` expected + it "returns false for for-statement" $ do + let stmt = For (nameTrg "x") (list IntTy []) [Return (constIntExp 0)] + let expected = False + doesAlwaysReturn stmt `shouldBe` expected diff --git a/test/Jikka/RestrictedPython/Language/VariableAnalysisSpec.hs b/test/Jikka/RestrictedPython/Language/VariableAnalysisSpec.hs new file mode 100644 --- /dev/null +++ b/test/Jikka/RestrictedPython/Language/VariableAnalysisSpec.hs @@ -0,0 +1,32 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Jikka.RestrictedPython.Language.VariableAnalysisSpec + ( spec, + ) +where + +import Jikka.RestrictedPython.Language.Expr +import Jikka.RestrictedPython.Language.VariableAnalysis +import Jikka.RestrictedPython.Language.WithoutLoc +import Test.Hspec + +spec :: Spec +spec = do + describe "analyzeStatementMax" $ do + it "works" $ do + let e = AnnAssign (nameTrg "y") IntTy (name "x") + let expected = (ReadList ["x"], WriteList ["y"]) + analyzeStatementMax e `shouldBe` expected + it "works'" $ do + let e = If (constBoolExp True) [AnnAssign (nameTrg "y") IntTy (name "x")] [] + let expected = (ReadList ["x"], WriteList ["y"]) + analyzeStatementMax e `shouldBe` expected + describe "analyzeStatementMin" $ do + it "works" $ do + let e = AnnAssign (nameTrg "y") IntTy (name "x") + let expected = (ReadList ["x"], WriteList ["y"]) + analyzeStatementMin e `shouldBe` expected + it "works'" $ do + let e = If (constBoolExp True) [AnnAssign (nameTrg "y") IntTy (name "x")] [] + let expected = (ReadList [], WriteList []) + analyzeStatementMin e `shouldBe` expected diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,1 @@ +{-# OPTIONS_GHC -F -pgmF hspec-discover #-}