Submission #286413

---

Source Code Expand

{-# LANGUAGE CPP #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverloadedStrings #-}

import qualified Data.ByteString.Char8 as S
import Data.Monoid

main :: IO ()
main = do
  (S.lines -> [S.words -> [S.unpack -> read -> (i0 :: Int)], S.words -> [S.unpack -> read -> (i1 :: Int)]]) <- S.getContents
  S.putStrLn $ if 1 < i0 `div` i1 then "YES" else "NO"

-- This program was generated and submitted by the following program:

# if 0

{-# LANGUAGE TupleSections #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE QuasiQuotes #-}

-- required packages (apart from the Haskell Platform): 
--  vector-space, tagsoup, lens, interpolatedstring-perl6, curl, http-types

module Main (main, logIn) where

import Control.Applicative
import Control.Concurrent
import Control.Lens hiding (argument, op, act)
import Control.Monad
import Control.Monad.Reader
import Control.Monad.State.Strict
import Data.AdditiveGroup
import qualified Data.ByteString.Char8 as S
import Data.Fixed
import Data.Function
import qualified Data.IntSet as IS
import Data.List
import qualified Data.Map as M
import Data.Maybe
import Data.Monoid
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.Text.IO as T
import qualified Data.Text.Lazy as TL
import qualified Data.Text.Lazy.Builder as TB
import qualified Data.Text.Lazy.Builder.Int as TB
import qualified Data.Text.Read as T
import qualified Data.Time as Time
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import Debug.Trace
import Language.Haskell.TH
import qualified Network.Curl as Curl
import qualified Network.HTTP.Types.URI as HTTPTypes
import qualified Text.HTML.TagSoup as TS
import qualified Text.HTML.TagSoup.Match as TS
import Text.InterpolatedString.Perl6

main :: IO ()
main = do
  sleepUntil $ utc 2014 11 29 12 00 03.0

  solveAndSubmit "arc030"

-- | Create a UTCTime.
utc :: Integer -> Int -> Int -> Int -> Int -> Pico -> Time.UTCTime
utc year month day hour minute sec =
  Time.UTCTime (Time.fromGregorian year month day) $
    Time.timeOfDayToTime $ Time.TimeOfDay hour minute sec

-- | Sleep until the specified time.
sleepUntil :: Time.UTCTime -> IO ()
sleepUntil !end = do
  now <- Time.getCurrentTime
  when (now < end) $ do
    let toSleep = 0.9 * realToFrac (Time.diffUTCTime end now) :: Double
    when (toSleep > 0.1) $ threadDelay $ floor $ toSleep * 1000000
    sleepUntil end

-- | Solve the first problem of the given contest and submit the solution.
solveAndSubmit :: T.Text -> IO ()
solveAndSubmit contest = do
  curl <- newCurl
  descr <- getProblemDescription curl contest
  let prog = solve descr
  T.writeFile "submitted.hs" prog
  submitProgram curl contest Haskell prog

-- | Given a problem description, create a Haskell source file that solves it.
solve :: T.Text -> T.Text
solve desc =
  makeProblem (extractInstances desc) (extractIntLiterals desc) $ \problem ->
    prettyProgram problem $ findSolution problem

----------------------------------------------------------------
-- I/O with AtCoder

-- | Create a new 'Curl' with default options set.
newCurl :: IO Curl.Curl
newCurl = do
  curl <- Curl.initialize
  Curl.setopts curl
    [ Curl.CurlCookieFile "cookies"
    , Curl.CurlCookieJar "cookies"
    , Curl.CurlUserAgent "AlmostBruteforceCoder/0.2"
    , Curl.CurlVerbose True
    ]
  return curl

-- | Get the HTML page describing the first problem of the given contest.
getProblemDescription :: Curl.Curl -> T.Text -> IO T.Text
getProblemDescription curl contest = do
  (200, content) <- getURL curl $ firstProblemPage contest
  return content

data Language = Perl | Haskell
  deriving (Show)

-- | Submit the given text as a solution to the first problem of the given
-- contest.
submitProgram :: Curl.Curl -> T.Text -> Language -> T.Text -> IO ()
submitProgram curl contest lang code = do
  (200, submitPageContents) <- getURL curl $ submitPage contest
  let !sinfo = parseSubmitInfo submitPageContents
  _ <- submitWithInfo curl contest sinfo lang code
  return ()

type URL = String
type Response = Curl.CurlResponse_ [(String, String)] S.ByteString

submitPage :: T.Text -> URL
submitPage contestName = [qc|http://{contestName}.contest.atcoder.jp/submit|]

submitPost :: T.Text -> T.Text -> URL
submitPost contestName firstPid =
  [qc|http://{contestName}.contest.atcoder.jp/submit?task_id={firstPid}|]

firstProblemPage :: T.Text -> URL
firstProblemPage contestName =
  [qc|http://{contestName}.contest.atcoder.jp/tasks/{contestName}_1|]

_loginPage :: URL
_loginPage = "http://abc010.contest.atcoder.jp/login?next_url=http%3A%2F%2Fabc010.contest.atcoder.jp%2Fsubmit#40160"

loginPost :: URL
loginPost = "http://abc010.contest.atcoder.jp/login?next_url=http%3A%2F%2Fabc010.contest.atcoder.jp%2Fsubmit"

logIn :: Curl.Curl -> IO (Int, T.Text)
logIn curl = do
  [atcoderId, atcoderPassword] <- T.lines <$> readUTF8File "atcoder_passwd"
  postURL curl loginPost
    [ ("name", atcoderId)
    , ("password", atcoderPassword)
    ]

-- | Read a UTF-8 encoded file.
readUTF8File :: FilePath -> IO T.Text
readUTF8File path = T.decodeUtf8 <$> S.readFile path

submitWithInfo
  :: Curl.Curl -> T.Text -> SubmitInfo -> Language -> T.Text
  -> IO (Int, T.Text)
submitWithInfo curl contest sinfo lang code =
    postURL curl (submitPost contest taskId) $
  [ ("__session", sessionId sinfo)
  , ("task_id", taskId)
  , ("source_code", code)
  ] ++ do
    tid <- problemIds sinfo
    return ("language_id_" <> tid, langId)
  where
    taskId = head' "no task" $ problemIds sinfo 
    langId = case lang of
      Perl -> "8"
      Haskell -> "11"

-- | Send a GET request to the given url, and return the status code and the
-- content of the page.
getURL :: Curl.Curl -> URL -> IO (Int, T.Text)
getURL curl url =
  checkResp =<< Curl.do_curl_ curl url [Curl.CurlHttpGet True]

-- | Send a POST request to the given url with the given set of parameters.
postURL :: Curl.Curl -> URL -> [(T.Text, T.Text)] -> IO (Int, T.Text)
postURL curl url params =
  checkResp =<< Curl.do_curl_ curl url [Curl.CurlPostFields pairs]
  where
    pairs = do
      (key, value) <- params
      return [qc|{enc key}={enc value}|]

    enc = HTTPTypes.urlEncode True . T.encodeUtf8

checkResp :: Response -> IO (Int, T.Text)
checkResp resp
  | Curl.respCurlCode resp /= Curl.CurlOK =
      fail [qc|bad curl status: {Curl.respCurlCode resp}|]
  | otherwise = return (Curl.respStatus resp, T.decodeUtf8 $ Curl.respBody resp)

-- | Information needed to make a POST request for submission.
data SubmitInfo = SubmitInfo
  { problemIds :: [T.Text]
  , sessionId :: T.Text
  }
  deriving (Show)

-- | Parse the submission page to extract a @SubmitInfo@.
parseSubmitInfo :: T.Text -> SubmitInfo
parseSubmitInfo txt@(TS.parseTags -> tags) = SubmitInfo
  { problemIds = pids
  , sessionId = sid
  }
  where
    sid = head' [qc|no session id found: {txt}|] $ do
      TS.TagOpen "input" attrs <- tags
      guard $ elem ("name", "__session") attrs
      ("value", value) <- attrs
      return value
    pids = head' [qc|no task selector found: {txt}|] $ do
      sec <-
        takeWhile (not . TS.tagCloseLit "select") <$>
        TS.sections (TS.tagOpenAttrLit "select" ("name", "task_id")) tags
      return $ do
        TS.TagOpen "option" attrs <- sec
        ("value", value) <- attrs
        return value

-- | 'head' with a better error message.
head' :: String -> [a] -> a
head' msg [] = error [qc|head': {msg}|]
head' _ (x:_) = x

----------------------------------------------------------------
-- Pretty-printing a program

-- | Create Haskell source text from a 'Program'.
prettyProgram :: Problem oty -> Program oty -> T.Text
prettyProgram prob prog = TL.toStrict $ TB.toLazyText builder
  where
    builder
      = TB.fromText prolog
      <> pprInputStatement (inputPattern $ problemSpec prob)
      <> "\n  "
      <> pprOutputStatement prog
      <> TB.fromText epilog

prolog :: T.Text
prolog = [q|
{-# LANGUAGE CPP #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverloadedStrings #-}

import qualified Data.ByteString.Char8 as S
import Data.Monoid

main :: IO ()
main = do
  |]

epilog :: T.Text
epilog = [qc|

-- This program was generated and submitted by the following program:

# if 0

{thisProgramText}
# endif
|]

-- | The source text of this program.
thisProgramText :: String
thisProgramText = $(litE . stringL =<< runIO (readFile "abc.hs"))

-- | Format an input pattern into a statement that binds input variables.
pprInputStatement :: InputPattern -> TB.Builder
pprInputStatement ip = pat <> " <- S.getContents"
  where
    pat = evalState (pprInputPatten ip) $ PatPprState 0 0

-- | Monad for pretty-printing an input pattern.
type PatPpr = State PatPprState

data PatPprState = PatPprState
  { _usedInts :: !Int -- ^ How many int variables have we bound?
  , _usedStrs :: !Int -- ^ Ditto for string variables
  }

usedInts :: Lens' PatPprState Int
usedInts = lens _usedInts $ \s x -> s{ _usedInts = x }

usedStrs :: Lens' PatPprState Int
usedStrs = lens _usedStrs $ \s x -> s{ _usedStrs = x }

pprInputPatten :: InputPattern -> PatPpr TB.Builder
pprInputPatten GetContents = strVarPattern
pprInputPatten (Lines ls) = do
  lps <- mapM pprLinePattern ls
  return $ "(S.lines -> " <> pprList lps <> ")"

pprLinePattern :: LinePattern -> PatPpr TB.Builder
pprLinePattern GetLine = strVarPattern
pprLinePattern (Words ws) = do
  wps <- mapM pprWordPattern ws
  return $ "S.words -> " <> pprList wps

pprWordPattern :: WordPattern -> PatPpr TB.Builder
pprWordPattern Id = strVarPattern
pprWordPattern Read = do
  varPat <- intVarPattern
  return $ "S.unpack -> read -> (" <> varPat <> " :: Int)"

pprList :: [TB.Builder] -> TB.Builder
pprList list = "[" <> go list
  where
    go [] = "]"
    go (x:xs) = x <> go1 xs

    go1 [] = "]"
    go1 (x:xs) = ", " <> x <> go1 xs

strVarPattern :: PatPpr TB.Builder
strVarPattern = do
  k <- use usedStrs
  usedStrs .= k + 1
  return $ pprStrArgument k

intVarPattern :: PatPpr TB.Builder
intVarPattern = do
  k <- use usedInts
  usedInts .= k + 1
  return $ pprIntArgument k

-- | Format a program into a statement that prints the output.
pprOutputStatement :: Program oty -> TB.Builder
pprOutputStatement (IntProgram ip)
  = (const "print" `apply` pprInt ip) (-100)
pprOutputStatement (BoolProgram ip)
  = "S.putStrLn $ if " <> pprBool ip (-100) <> " then \"YES\" else \"NO\""
pprOutputStatement (StrProgram ip)
  = (const "S.putStrLn" `apply` pprStr ip) (-100)

pprInt :: IntProgram -> Int -> TB.Builder
pprInt prog = case prog of
  a :+ b -> parens 6 $ pprInt a 6 <> " + " <> pprInt b 7
  a :- b -> parens 6 $ pprInt a 6 <> " - " <> pprInt b 7
  a :* b -> parens 7 $ pprInt a 7 <> " * " <> pprInt b 8
  a :/ b -> parens 7 $ pprInt a 7 <> " `div` " <> pprInt b 8
  a :/^ b -> const "(\\x y -> (x + y - 1) `div` y)"
    `apply` pprInt a `apply` pprInt b
  a :% b -> parens 7 $ pprInt a 7 <> " `mod` " <> pprInt b 8
  Max a b -> const "max" `apply` pprInt a `apply` pprInt b
  Min a b -> const "min" `apply` pprInt a `apply` pprInt b
  NumChars s -> const "S.length" `apply` pprStr s
  NumLines s -> const "((+1) . S.count '\n')" `apply` pprStr s
  LiteralInt n -> const $ TB.decimal n
  IntArgument k -> const $ pprIntArgument k

pprStr :: StrProgram -> Int -> TB.Builder
pprStr prog = case prog of
  a :++ b -> parens 6 $ pprStr a 7 <> " <> " <> pprStr b 6
  Show a -> const "S.pack" `apply` (const "show" `apply` pprInt a)
  Sep a b -> parens 6 $ pprStr a 7 <> " <> \" \" <> "  <> pprStr b 6
  VCat a b -> parens 6 $ pprStr a 7 <> " <> \"\\n\" <> "  <> pprStr b 6
  LiteralString str -> const $ TB.fromString (show str)
  StrArgument k -> const $ pprStrArgument k

pprBool :: BoolProgram -> Int -> TB.Builder
pprBool prog = case prog of
  a :&& b -> parens 3 $ pprBool a 4 <> " && " <> pprBool b 3
  a :|| b -> parens 2 $ pprBool a 3 <> " && " <> pprBool b 2
  Not a -> const "not" `apply` pprBool a
  a :< b -> parens 4 $ pprInt a 5 <> " < " <> pprInt b 5
  a :== b -> parens 4 $ pprInt a 5 <> " == " <> pprInt b 5

pprIntArgument :: Int -> TB.Builder
pprIntArgument k = "i" <> TB.decimal k

pprStrArgument :: Int -> TB.Builder
pprStrArgument k = "s" <> TB.decimal k

apply :: (Int -> TB.Builder) -> (Int -> TB.Builder) -> Int -> TB.Builder
apply f x = parens 10 $ f 10 <> " " <> x 11

parens :: Int -> TB.Builder -> Int -> TB.Builder
parens k b prec
  | prec > k = "(" <> b <> ")"
  | otherwise = b

----------------------------------------------------------------
-- Finding a solution

-- | Return a program that passes all the test cases. If there is no such
-- representable program, it loops forever.
findSolution :: (ObjectType oty) => Problem oty -> Program oty
findSolution prob = head $ filter good $ generateAll prob
  where
    good prog = all (satisfiesExample prog) $ instances prob

-- | Returns an inifite list that enumerates (almost) all representable
-- programs.
generateAll :: (ObjectType oty) => Problem oty -> [Program oty]
generateAll prob = concatMap gen [0..]
  where
    spec = problemSpec prob
    ints = intsInSpec prob
    gen sz = generateProgramsWithSize spec ints sz
      ++ do
        guard $ sz' >= 0
        (prob', embed) <- reds
        map embed $ generateProgramsWithSize (problemSpec prob') ints sz'
      where
        sz' = sz - 2
    reds = reductions prob

-- | Test whether the given program behaves in the way that is specified
-- in an example.
satisfiesExample :: (ObjectType oty) => Program oty -> TypedInstance oty -> Bool
satisfiesExample prog (TypedInstance inputs expected) =
  eval prog inputs == Just expected

-- | Evaluate the given program with the given inputs. Returns the outcmome
-- of the evaluation, or @Nothing@ if an error occurs.
eval :: Program oty -> TypedInput -> Maybe oty
eval prog input = runReaderT go input
  where
    go = case prog of
      IntProgram p -> evalInt p
      StrProgram p -> evalString p
      BoolProgram p -> evalBool p

-- | A monad for evaluating a program.
type Eval = ReaderT TypedInput Maybe

evalString :: StrProgram -> Eval S.ByteString
evalString prog = seqM $ case prog of
  a :++ b -> (<>) <$> evalString a <*> evalString b
  Sep a b -> do
    x <- evalString a
    y <- evalString b
    return $! x <> " " <> y
  VCat a b -> do
    x <- evalString a
    y <- evalString b
    return $! x <> "\n" <> y
  Show a -> S.pack . show <$> evalInt a
  LiteralString a -> return a
  StrArgument k -> asks $ (V.!k) . strInput

evalInt :: IntProgram -> Eval Int
evalInt prog = seqM $ case prog of
  a :+ b -> (+) <$> evalInt a <*> evalInt b
  a :- b -> (-) <$> evalInt a <*> evalInt b
  a :* b -> (*) <$> evalInt a <*> evalInt b
  a :/ b -> div <$> evalInt a <*> (nonzero =<< evalInt b)
  a :/^ b -> cdiv <$> evalInt a <*> (nonzero =<< evalInt b)
  a :% b -> mod <$> evalInt a <*> (nonzero =<< evalInt b)
  Max a b -> max <$> evalInt a <*> evalInt b
  Min a b -> min <$> evalInt a <*> evalInt b
  NumChars a -> S.length <$> evalString a
  NumLines a -> (+1) . S.count '\n' <$> evalString a
  LiteralInt n -> return n
  IntArgument k -> asks $ (U.!k) . intInput
  where
    cdiv a b = (a+b-1) `div` b

    nonzero 0 = mzero
    nonzero x = return x

evalBool :: BoolProgram -> Eval Bool
evalBool prog = seqM $ case prog of
  a :&& b -> ifM (evalBool a) (evalBool b) (return False)
  a :|| b -> ifM (evalBool a) (return True) (evalBool b)
  Not a -> not <$> evalBool a
  a :< b -> (<) <$> evalInt a <*> evalInt b
  a :== b -> (==) <$> evalInt a <*> evalInt b

seqM :: (Monad m) => m a -> m a
seqM = (>>= (return$!))
{-# INLINE seqM #-}

ifM :: (Monad m) => m Bool -> m a -> m a -> m a
ifM a t f = a >>= \case
  True -> t
  False -> f
{-# INLINE ifM #-}

----------------------------------------------------------------
-- Reducing a problem

-- | Try to reduce the problem a bit by looking for a common pattern
-- in example outputs. Returns a list of possible reductions. Each reduction
-- is a pair: the first element is the reduced
-- problem and the second element is a function that takes a solution
-- to the reduced problem and produces a solution to the original problem.
--
-- For example, if all the sample outputs end with a "pp" suffix, it generates
-- a "reduced" problem whose task is to generate the output string without the
-- "pp" suffix.
reductions
  :: forall oty
  . (ObjectType oty)
  => Problem oty
  -> [(Problem oty, Program oty -> Program oty)]
reductions prob = do
  modifier <- modifiers
  (outputs, embed) <- modifier $ map expectedOutput iss
  let newProb = prob{ instances = zipWith replaceOutput iss outputs }
  return (newProb, embed)
  where
    iss = instances prob
    replaceOutput inst out = inst{ expectedOutput = out }

    modifiers :: [[oty] -> [([oty], Program oty -> Program oty)]]
    modifiers = case objectType :: ObjectTypeTag oty of
      Bool -> []
      String -> [prefix, suffix]
      Int -> [add, mul]

    prefix = strip commonPrefixes S.drop (:++)
    suffix = strip commonSuffixes (\n s -> S.take (S.length s-n) s) (flip (:++))

    add, mul :: [Int] -> [([Int], Program Int -> Program Int)]
    add xs = do
      guard $ m /= 0
      let
        embed (IntProgram s) = IntProgram (s :+ LiteralInt m)
      return (map (subtract m) xs, embed)
      where
        m = minimum xs

    mul xs = do
      guard $ g /= 0 && g /= 1
      let
        embed (IntProgram s) = IntProgram (s :* LiteralInt g)
      return (map (flip div g) xs, embed)
      where
        g = foldl1' gcd xs

    strip
      :: ([S.ByteString] -> [S.ByteString])
      -> (Int -> S.ByteString -> S.ByteString)
      -> (StrProgram -> StrProgram -> StrProgram)
      -> [S.ByteString]
      -> [([S.ByteString], Program S.ByteString -> Program S.ByteString)]
    strip common f g (ss :: [S.ByteString]) = do
      c <- common ss
      guard $ not $ S.null c
      let
        embed (StrProgram s) = StrProgram (g (LiteralString c) s)
      return (map (f $ S.length c) ss, embed)

-- | Returns the list of all strings that are a prefix of all given strings.
--
-- When the empty list is given, returns @[""]@.
commonPrefixes :: [S.ByteString] -> [S.ByteString]
commonPrefixes = S.inits . commonPrefix

commonPrefix :: [S.ByteString] -> S.ByteString
commonPrefix [] = S.empty
commonPrefix xs = foldl1' common xs
  where
    common x y = S.take len x
      where
        len = length $ takeWhile id $ S.zipWith (==) x y

commonSuffixes :: [S.ByteString] -> [S.ByteString]
commonSuffixes = map S.reverse . commonPrefixes . map S.reverse

----------------------------------------------------------------
-- Generating programs

-- | A program that produces a value of type @oty@.
data Program oty where
  IntProgram :: IntProgram -> Program Int
  StrProgram :: StrProgram -> Program S.ByteString
  BoolProgram :: BoolProgram -> Program Bool

deriving instance Show (Program oty)

-- | A program that returns an Int.
data IntProgram
  = !IntProgram :+ !IntProgram
  | !IntProgram :- !IntProgram
  | !IntProgram :* !IntProgram
  | !IntProgram :/ !IntProgram -- ^ floor after division
  | !IntProgram :/^ !IntProgram -- ^ ceiling after division
  | !IntProgram :% !IntProgram
  | Max !IntProgram !IntProgram
  | Min !IntProgram !IntProgram
  | NumChars !StrProgram
  | NumLines !StrProgram
  | LiteralInt !Int
  | IntArgument !Int
  deriving (Show)

-- | A program that returns a Bool.
data BoolProgram
  = !BoolProgram :&& !BoolProgram
  | !BoolProgram :|| !BoolProgram
  | Not !BoolProgram
  | !IntProgram :< !IntProgram
  | !IntProgram :== !IntProgram
  deriving (Show)

-- | A program that returns a String.
data StrProgram
  = !StrProgram :++ !StrProgram -- ^ concatenation
  | Sep !StrProgram !StrProgram -- ^ concatenation with a space
  | VCat !StrProgram !StrProgram -- ^ concatenation with a newline
  | Show !IntProgram
  | LiteralString !S.ByteString
  | StrArgument !Int
  deriving (Show)

-- | A program-generation monad.
type Generate = ReaderT GenEnv (StateT GenState [])

-- | An argument to a program-generation process.
data GenEnv = GenEnv
  { _targetProblem :: !ProblemSpec
  , _size :: !Int
  , _allowedIntLiterals :: ![Int]
  }

targetProblem :: Lens' GenEnv ProblemSpec
targetProblem = lens _targetProblem $ \s x -> s{ _targetProblem = x }

size :: Lens' GenEnv Int
size = lens _size $ \s x -> s{ _size = x }

allowedIntLiterals :: Lens' GenEnv [Int]
allowedIntLiterals = lens _allowedIntLiterals $ \s x -> s{ _allowedIntLiterals = x }

-- | A state of a program-generation process.
data GenState = GenState
  { _intVars :: !IS.IntSet -- ^ set of used int inputs
  , _strVars :: !IS.IntSet -- ^ set of used string inputs
  }

intVars :: Lens' GenState IS.IntSet
intVars = lens _intVars $ \s x -> s{ _intVars = x }

strVars :: Lens' GenState IS.IntSet
strVars = lens _strVars $ \s x -> s{ _strVars = x }

generateProgramsWithSize
  :: forall oty. (ObjectType oty)
  => ProblemSpec -> IS.IntSet -> Int -> [Program oty]
generateProgramsWithSize spec ints sz = case objectType :: ObjectTypeTag oty of
  Int -> map IntProgram $ gen1 intProgram
  String -> map StrProgram $ gen1 strProgram
  Bool -> map BoolProgram $ gen1 boolProgram
  where
    gen1 :: Generate a -> [a]
    gen1 g = runGenerate spec ints sz g

-- | Run the given program generator to generate programs of the given size.
runGenerate :: ProblemSpec -> IS.IntSet -> Int -> Generate a -> [a]
runGenerate prob ints sz
  = map fst
  . filter (full . snd)
  . flip runStateT initialState
  . flip runReaderT env
  where
    env = GenEnv prob sz lits
    initialState = GenState IS.empty IS.empty
    lits = IS.toList $ IS.delete 0 $ IS.fromList $ do
      int <- IS.toList ints
      [int, int + 1]

    -- Filter out programs that don't use all the input variables.
    full st
      = st^.intVars.to IS.size == nIntInputs prob
      && st^.strVars.to IS.size == nStringInputs prob

-- | Generate @IntProgram@s.
intProgram :: Generate IntProgram
intProgram = msum
  [ binIntOpSym (:+)
  , binIntOp (:-)
  , binIntOpSym (:*)
  , binIntOp (:/)
  , binIntOp (:/^)
  , binIntOp (:%)
  , binIntOpSym Min
  , binIntOpSym Max
  , unOp NumChars strProgram
  , unOp NumLines strProgram
  , literal
  , argument nIntInputs intVars IntArgument
  ]
  where
    binIntOp op = binOp op intProgram intProgram
    binIntOpSym op = binOpSym op intProgram intProgram
    literal = do
      1 <- view size
      ints <- view allowedIntLiterals
      msum $ map (return . LiteralInt) ints

boolProgram :: Generate BoolProgram
boolProgram = msum
  [ binOpSym (:&&) boolProgram boolProgram
  , binOpSym (:||) boolProgram boolProgram
  , unOp Not boolProgram
  , binOp (:<) intProgram intProgram
  , binOpSym (:==) intProgram intProgram
  ]

strProgram :: Generate StrProgram
strProgram = msum
  [ binOp (:++) strProgram strProgram
  , binOp Sep strProgram strProgram
  , binOp VCat strProgram strProgram
  , unOp Show intProgram
  , argument nStringInputs strVars StrArgument
  ]

argument
  :: (ProblemSpec -> Int)
  -> Lens' GenState IS.IntSet
  -> (Int -> a)
  -> Generate a
argument inputCount vars f = do
  sz <- view size
  guard $ sz <= 2
  nvars <- view $ targetProblem.to inputCount
  varId <- msum $ map return [0..nvars-1]
  used <- IS.member varId <$> use vars

  -- Impose a bigger cost when using an already-used variable.
  guard $ sz == if used then 2 else 1
  vars %= IS.insert varId
  return $ f varId

unOp :: (a -> b) -> Generate a -> Generate b
unOp op genChild = do
  sz <- view size
  guard $ 2 <= sz
  local (size .~ sz - 1) $ op <$> genChild

binOpSym :: (a -> b -> c) -> Generate a -> Generate b -> Generate c
binOpSym = binOpMaybeSym True

binOp :: (a -> b -> c) -> Generate a -> Generate b -> Generate c
binOp = binOpMaybeSym False

binOpMaybeSym :: Bool -> (a -> b -> c) -> Generate a -> Generate b -> Generate c
binOpMaybeSym symmetric op leftChild rightChild = do
  !sz <- view size
  guard $ 3 <= sz
  let !ub = if symmetric then div sz 2 else sz - 2
  !leftSize <- msum $ map return [1 .. ub]
  let !rightSize = sz - leftSize - 1
  op
    <$> local (size .~ leftSize) leftChild
    <*> local (size .~ rightSize) rightChild

----------------------------------------------------------------
-- Guessing Input/Output formats of a problem

-- | A problem with a defined I/O format.
data Problem otype = Problem
  { problemSpec :: !ProblemSpec
  , instances :: ![TypedInstance otype]
  , intsInSpec :: !IS.IntSet -- ^ integer literals found in the spec
  }
  deriving (Show)

-- | The I/O specification of a problem, i.e. the number of arguments,
-- the type of each argument, etc.
data ProblemSpec = ProblemSpec
  { nIntInputs :: !Int
  , nStringInputs :: !Int
  , inputPattern :: !InputPattern
  , outputPattern :: !OutputPattern
  }

deriving instance Show ProblemSpec

-- | A problem instance in the parsed form.
data TypedInstance otype = TypedInstance
  { typedInput :: !TypedInput
  , expectedOutput :: !otype
  }
  deriving (Show)

-- | A concrete input to a problem.
data TypedInput = TypedInput
  { intInput :: !(U.Vector Int)
  , strInput :: !(V.Vector S.ByteString)
  }
  deriving (Show)

-- | A type of values that can be dealt with our program. Currently
-- it's one of Int, Bool or String.
class (Show a, Eq a) => ObjectType a where
  objectType :: ObjectTypeTag a

data ObjectTypeTag a where
  Int :: ObjectTypeTag Int
  String :: ObjectTypeTag S.ByteString
  Bool :: ObjectTypeTag Bool

deriving instance Show (ObjectTypeTag a)
deriving instance Eq (ObjectTypeTag a)

instance ObjectType Int where
  objectType = Int

instance ObjectType S.ByteString where
  objectType = String

instance ObjectType Bool where
  objectType = Bool

-- | How an input is encoded as a string, or an instruction to
-- systematically parse the input string into values.
--
-- Some patterns are more structured than others. For example,
-- any pattern of the form @(Lines x)@ is more structured than
-- @GetContents@.
data InputPattern
  = Lines [LinePattern] -- ^ Parse each line individually
  | GetContents -- ^ Take the whole input as a string
  deriving (Show)

-- | Instruction to parse a line.
data LinePattern
  = Words [WordPattern] -- ^ Parse each space-separated word individually
  | GetLine -- ^ Take the whole line as a string
  deriving (Show)

-- | Instruction to parse a space-separated word.
data WordPattern
  = Read -- ^ Read the word as an integer
  | Id -- ^ Take the word as a string
  deriving (Show)

-- | Instruction to parse an output example. Let's assume that the output
-- always consists of a single value
data OutputPattern = forall a. (ObjectType a) =>
  OutputPattern (ObjectTypeTag a)

deriving instance Show OutputPattern
instance Eq OutputPattern where
  OutputPattern Int == OutputPattern Int = True
  OutputPattern Bool == OutputPattern Bool = True
  OutputPattern String == OutputPattern String = True
  _ == _ = False

-- | Create a problem from examples.
makeProblem
  :: [Instance]
  -> IS.IntSet
  -> (forall oty. (ObjectType oty) => Problem oty -> r)
  -> r
makeProblem is ints cont =
  parseInstances (inputPattern spec) (outputPattern spec) is $ \ts ->
    cont $ Problem spec ts ints
  where
    spec = guessProblemSpec is

-- | Guess the I/O specificaitons using examples.
guessProblemSpec :: [Instance] -> ProblemSpec
guessProblemSpec is = ProblemSpec{..}
  where
    (nIntInputs, nStringInputs) = countInputs inputPattern
    inputPattern = lubs "guessInputPattern: no instance" $
      map (bestInputPattern . fst) is
    outputPattern = lubs "guessInputPattern: no instance" $
      map (bestOutputPattern . snd) is

-- | Parse the given instances using the given I/O patterns. This function is
-- in the CPS in order to avoid an explicit use of existential types.
parseInstances
  :: InputPattern
  -> OutputPattern
  -> [Instance]
  -> (forall otype. (ObjectType otype) => [TypedInstance otype] -> r)
  -> r
parseInstances ipat (OutputPattern (_ :: ObjectTypeTag oty)) is cont
  = cont (map (parseInstance ipat) is :: [TypedInstance oty])

-- | Parse one problem instance according to the input pattern.
-- The output pattern is inferred from the type.
parseInstance
  :: (ObjectType otype)
  => InputPattern -> Instance -> TypedInstance otype
parseInstance ipat (input, output) = TypedInstance
  { typedInput = TypedInput
    { intInput = U.fromList ints
    , strInput = V.fromList strs
    }
  , expectedOutput = parseOutput output
  }
  where
    (ints, strs) = parseInput ipat input

-- | Parse an input according to the input pattern.
parseInput :: InputPattern -> T.Text -> ([Int], [S.ByteString])
parseInput (Lines ps) txt = mconcat $ zipWith parseLine ps $ T.lines txt
parseInput GetContents txt = ([], [readString txt])

parseLine :: LinePattern -> T.Text -> ([Int], [S.ByteString])
parseLine (Words ps) txt = mconcat $ zipWith parseWord ps $ T.words txt
parseLine GetLine txt = ([], [readString txt])

parseWord :: WordPattern -> T.Text -> ([Int], [S.ByteString])
parseWord Read t = ([readInt t], [])
parseWord Id t = ([], [readString t])

parseOutput :: forall otype. (ObjectType otype) => T.Text -> otype
parseOutput txt = case objectType :: ObjectTypeTag otype of
  Int -> readInt txt
  String -> readString txt
  Bool -> fromMaybe (error [qc|cannot read bool: {show txt}|]) $
    readBoolMaybe txt

readString :: T.Text -> S.ByteString
readString = T.encodeUtf8

readInt :: T.Text -> Int
readInt txt = fromMaybe (error [qc|cannot read int: {show txt}|]) $
  readIntMaybe txt

readIntMaybe :: T.Text -> Maybe Int
readIntMaybe txt = case T.signed T.decimal txt of
  Right (val, rest)
    | T.null rest -> Just val
  _ -> Nothing

readBoolMaybe :: T.Text -> Maybe Bool
readBoolMaybe "YES" = Just True
readBoolMaybe "NO" = Just False
readBoolMaybe _ = Nothing

-- | Count the number of integer and string inputs in a pattern.
countInputs :: InputPattern -> (Int, Int)
countInputs (Lines ls) = sumV $ map countInputsInLine ls
countInputs GetContents = (0, 1)

countInputsInLine :: LinePattern -> (Int, Int)
countInputsInLine (Words ws) = sumV $ map countInputsInWord ws
countInputsInLine GetLine = (0, 1)

countInputsInWord :: WordPattern -> (Int, Int)
countInputsInWord Read = (1, 0)
countInputsInWord Id = (0, 1)

-- | Semilattice, or commutative idempotent semigroup.
class Semilattice a where
  lub :: a -> a -> a

-- The {Input,Line,Word}Pattern types each forms a semilattice where
-- @(a `lub` b)@ means the "most structured" pattern that is compatible
-- with both @a@ and @b@.
instance Semilattice InputPattern where
  lub (Lines xs) (Lines ys) =
    maybe GetContents Lines $ lubList xs ys
  lub _ _ = GetContents

instance Semilattice LinePattern where
  lub (Words xs) (Words ys) = 
    maybe GetLine Words $ lubList xs ys
  lub _ _ = GetLine

instance Semilattice WordPattern where
  lub Read Read = Read
  lub _ _ = Id

instance Semilattice OutputPattern where
  lub a b
    | a == b = a
  lub _ _ = OutputPattern String

-- | Take the lub of the given values.
lubs :: (Semilattice a) => String -> [a] -> a
lubs msg [] = error msg
lubs _ is = foldl1' lub is

-- | Pairwise lub, if the lists have the same length.
lubList :: (Semilattice a) => [a] -> [a] -> Maybe [a]
lubList xs ys
  | length xs == length ys = Just $ zipWith lub xs ys
  | otherwise = Nothing

-- | Return the most structured input pattern that can parse the
-- given text.
bestInputPattern :: T.Text -> InputPattern
bestInputPattern = Lines . map bestLinePattern . T.lines

bestLinePattern :: T.Text -> LinePattern
bestLinePattern = Words . map bestWordPattern . T.words

bestWordPattern :: T.Text -> WordPattern
bestWordPattern txt
  | isJust (readIntMaybe txt) = Read
  | otherwise = Id

bestOutputPattern :: T.Text -> OutputPattern
bestOutputPattern txt
  | isJust (readIntMaybe txt) = OutputPattern Int
  | isJust (readBoolMaybe txt) = OutputPattern Bool
  | otherwise = OutputPattern String

----------------------------------------------------------------
-- Extracting problem instances from HTML

-- | A pair of sample input and output.
type Instance = (T.Text, T.Text)

-- | Extract problem instances from a problem page.
extractInstances :: T.Text -> [Instance]
extractInstances
  = combineSamples
  . map extractSample
  . TS.partitions (isJust . parseSampleHead)
  . TS.parseTags

-- | A sample input or output.
data Sample = Sample
  { _inputOutput :: !InputOutput
  , _instanceName :: !T.Text
  , _sampleBody :: !T.Text
  }

data InputOutput = Input | Output
  deriving (Eq)

-- | Try to parse the given tag as an indicator of a sample.
parseSampleHead :: TS.Tag T.Text -> Maybe (InputOutput, T.Text)
parseSampleHead (TS.TagText (T.strip -> txt)) =
  ((Input,) <$> T.stripPrefix "入力例" txt) <|>
  ((Output,) <$> T.stripPrefix "出力例" txt)
parseSampleHead _ = Nothing

-- | Extract one sample from HTML.
extractSample :: [TS.Tag T.Text] -> Sample
extractSample
  ((parseSampleHead -> Just (io, name)) : ts0)
    = case TS.partitions (TS.tagOpenNameLit "pre") ts0 of
      [] -> error "extractSample: no pre found"
      ts1 : _ -> fromMaybe (error "extractSample: no text found") $
        fmap (Sample io name . T.strip) $ msum $ map TS.maybeTagText ts1
extractSample ts = error [qc|extractSample: no head: {take 1 ts}|]

-- | Combine matching input and output samples to build instances.
combineSamples :: [Sample] -> [Instance]
combineSamples samples
  = mapMaybe buildInstance
  $ M.toList
  $ M.fromListWith (++)
  [ (name, [(io, body)]) | Sample io name body <- samples]
  where
    buildInstance (name, ss) = case partition ((==Input) . fst) ss of
      ([(_, input)], [(_, output)]) -> Just (input, output)
      ([], _) -> warn "no input"
      (_, []) -> warn "no output"
      (_:_:_, _) -> warn "multiple inputs"
      (_, _:_:_) -> warn "multiple outputs"
      where
        warn msg = trace [qc|combineSamples: {name}: {msg::String}|] Nothing

----------------------------------------------------------------
-- Extracting integer literals from HTML

-- | Extract integer literals from a problem page.
extractIntLiterals :: T.Text -> IS.IntSet
extractIntLiterals
  = IS.fromList
  . concatMap extractIntsFromText
  . mapMaybe TS.maybeTagText
  . interestingRegion
  . TS.parseTags
  where
    interestingRegion
      = snd . break (TS.tagOpenAttrLit "div" ("id", "task-statement"))
      . fst . break (TS.tagText $ (=="入力例1") . T.strip)

extractIntsFromText :: T.Text -> [Int]
extractIntsFromText text = do
  grp <- T.groupBy ((==) `on` isAsciiDigit) text
  maybeToList $ readIntMaybe grp
  where
    isAsciiDigit x = '0' <= x && x <= '9'

# endif

Submission Info

Judge Result