prog_synthesis/escher/Main.hs

import Expr
import Eval
import Syntesis

import Control.Monad (guard, liftM, when, unless, foldM, foldM_)
import Control.Applicative
import Control.Monad.State as State
import Data.Map (Map)
import Data.Set (Set)
import qualified Data.Map as Map
import qualified Data.Set as Set
import Data.List (inits)
import Data.Maybe (fromMaybe, isJust, maybeToList, isNothing)

import Debug.Trace (trace)
import TypeCheck
import qualified Data.List as List


syntEval :: [Value] -> Expr -> SyntState (Result Value)
syntEval input expr = do cache <- gets syntCache
                         conf <- gets $ confBySynt input expr False
                         return $ eval conf expr

syntEvalEx :: [Value] -> Expr -> SyntState (Result Value)
syntEvalEx input expr = do cache <- gets syntCache
                           conf <- gets $ confBySynt input expr True
                           return $ eval conf expr

syntCacheEval :: [Value] -> Expr -> SyntState (Result Value)
syntCacheEval input expr = do cache <- gets syntCache
                              conf <- gets $ confBySynt input expr False
                              let (result, cache') = cachedEval cache conf expr
                              modify $ \st -> st {syntCache = cache'}
                              return result

------------

matchGoal :: Goal -> Expr -> SyntState Bool
matchGoal (Goal goal) expr = do examples <- gets syntExamples
                                foldM checkOnInput True $ zip examples goal
  where checkOnInput False _ = return False
        checkOnInput acc (input, output) = do output' <- syntEval input expr
                                              return $ matchValue output' output -- TODO
        matchValue (Result x) (Just y) = x == y
        matchValue _ Nothing = True
        matchValue _ _ = False

------ syntesis steps

calcExprOutputs :: Expr -> SyntState [Result Value]
calcExprOutputs expr = do examples <- gets syntExamples
                          mapM (`syntEval` expr) examples -- OR: syntCacheEval (slower?)

calcExprExOutputs :: Expr -> SyntState [Result Value]
calcExprExOutputs expr = do examples <- gets syntExamples
                            mapM (`syntEvalEx` expr) examples

calcTemporaryExprOutputs :: Expr -> SyntState [Result Value]
calcTemporaryExprOutputs expr = do examples <- gets syntExamples
                                   mapM (`syntEval` expr) examples

matchAnyOutputs :: [Result Value] -> SyntState Bool
matchAnyOutputs outputs = do exprs <- gets syntExprs
                             foldM step False $ map fst exprs
  where step :: Bool -> Expr -> SyntState Bool
        step True _ = return True
        step False expr = do exprOutputs <- calcExprExOutputs expr -- NOTE: compare with examples to not throw out elements that possibly diffenernt on new examples
                             return $ outputs == exprOutputs -- and $ zipWith sameResults outputs exprOutputs
        sameResults (Result left) (Result right) = left == right
        sameResults (RecError {}) (RecError {}) = True
        sameResults _ _ = False

-- generate next step of exprs, remove copies
forwardStep :: Expr -> [Expr] -> SyntState (Maybe Expr)
forwardStep comp args = do let expr = fillHoles comp args
                           typeConf <- gets $ oracleTypes . syntOracle
                           if isNothing $ checkType typeConf expr
                           then return Nothing
                           else do
                             outputs <- calcExprExOutputs expr -- NOTE: compare with examples to not throw out elements that possibly diffenernt on new examples
                             matchedExisting <- matchAnyOutputs outputs
                             -- TODO: all RecErrors example could be useful on future cases ?
                             if matchedExisting || any isFatalError outputs || all isRecError outputs
                             then return Nothing
                             else do
                               modify $ \st -> st { syntExprs = (expr, []) : syntExprs st}
                               return $ Just expr

splitGoal :: Goal -> [Bool] -> Resolver
splitGoal resolverGoal@(Goal outputs) selector | length outputs == length selector =
  let resolverCond = Goal $ map (Just . BoolV) selector in
  let resolverThen = Goal $ zipWith (\v b -> if b then v else Nothing) outputs selector in
  let resolverElse = Goal $ zipWith (\v b -> if b then Nothing else v) outputs selector in
  Resolver { resolverGoal, resolverCond, resolverThen, resolverElse }

-- split goal by its index and by expr (if any answers matched), check if there is same goals to generated
splitGoalStep :: Goal -> [Bool]  -> SyntState Resolver
splitGoalStep goal selector = do let r = splitGoal goal selector
                                 resolvers <- gets syntResolvers
                                 unless (r `elem` resolvers) $ -- do not add existing resolvers -- do not add existing resolvers -- do not add existing resolvers -- do not add existing resolvers
                                    -- do not add existing resolvers
                                    -- do not add existing resolvers
                                    -- do not add existing resolvers -- do not add existing resolvers
                                    -- do not add existing resolvers
                                   modify $ \st -> st { syntUnsolvedGoals = Set.insert (resolverCond r) $
                                                                            Set.insert (resolverThen r) $
                                                                            Set.insert (resolverElse r) $
                                                                            syntUnsolvedGoals st,
                                                        syntResolvers = r : syntResolvers st } -- Set.insert
                                 return r

-- TODO: use expr evaluated outputs ?
trySolveGoal :: Expr -> Goal -> SyntState Bool
trySolveGoal expr goal = do doesMatch <- matchGoal goal expr
                            if doesMatch then do
                              modify $ \st -> st { syntSolvedGoals = Map.insert goal expr $ syntSolvedGoals st --,
                                                   -- syntUnsolvedGoals = Set.delete goal $ syntUnsolvedGoals st
                                      }
                              return True
                              -- trace ("goal solved: " ++ show goal) -- Tmp: trace
                            else return False

isGoalSolved :: Goal -> SyntState Bool
isGoalSolved goal = gets (Map.member goal . syntSolvedGoals)

goalSolution :: Goal -> SyntState (Maybe Expr)
goalSolution goal = gets (Map.lookup goal . syntSolvedGoals)

-- find all goals solved by new expr, by expr id it's values on examples, remove solved goals
-- returns found expr
-- NOTE: goals expected to be resolved
resolveStep :: (Expr, Expr, Expr) -> Resolver -> SyntState Expr
-- resolveStep r _ | trace ("resolution: " ++ show r) False = undefined -- TMP: trace
resolveStep (ifCond, ifDoThen, ifDoElse) r = do let expr = IfE { ifCond, ifDoThen, ifDoElse }
                                                let goal = resolverGoal r
                                                modify $ \st -> st { syntSolvedGoals = Map.insert goal expr $ syntSolvedGoals st,
                                                                     -- syntUnsolvedGoals = Set.delete goal $ syntUnsolvedGoals st,
                                                                     syntExprs = (expr, []) : syntExprs st }
                                                return expr

tryResolve :: Resolver -> SyntState (Maybe Expr)
-- tryResolve r | trace ("try resolution: " ++ show r) False = undefined -- TMP: trace
tryResolve r = do condSol <- goalSolution $ resolverCond r
                  thenSol <- goalSolution $ resolverThen r
                  elseSol <- goalSolution $ resolverElse r
                  case (condSol, thenSol, elseSol) of
                    (Just condExpr, Just thenExpr, Just elseExpr) -> do
                      expr <- resolveStep (condExpr, thenExpr, elseExpr) r
                      return $ Just expr
                    _ -> return Nothing

remakeSynt :: [[Value]] -> [Value] -> SyntState ()
remakeSynt newInputs newOutputs = do st <- get
                                     let Goal oldOutputs = syntRoot st
                                     goals <- gets $ \st -> zip (newInputs ++ syntExamples st)
                                                                (newOutputs ++ map (fromMaybe undefined) oldOutputs)
                                     initSynt (syntOracle st) goals
                                     modify (\st' -> st' { syntExprs = syntExprs st })

-- clear goal tree up to root, add example, calculate exprs on input (could be recursive ?)
-- returns new example
saturateStep :: Expr -> SyntState [[Value]]
saturateStep expr = do examples <- gets syntExamples
                       (newInputs, newOutputs) <- unzip <$> foldM searchEx [] examples
                       let exFound = not . null $ newInputs
                       when exFound $ remakeSynt newInputs newOutputs
                       return newInputs
    where searchEx prevExs input = do output <- syntEvalEx input expr
                                      return $ List.union prevExs $ case output of -- union ??
                                        NewExamples exs -> exs
                                        _ -> []

-- try to find terminating expr
terminateStep :: SyntState (Maybe Expr)
terminateStep = do rootGoal <- gets syntRoot
                   -- TODO: move try solve goal there ??
                   gets $ Map.lookup rootGoal . syntSolvedGoals
                   -- NOTE: Goal should be already solved earlier
                   -- doesMatch <- matchGoal rootGoal expr
                   -- return $ if doesMatch then Just expr else Nothing

-- TODO: FIXME
tryTerminate :: Expr -> SyntState (Maybe Expr)
tryTerminate expr = do maybeResult <- terminateStep
                       if isJust maybeResult
                       then do
                         newEx <- saturateStep expr
                         if not . null $ newEx
                         then do
                           -- redo prev exprs (including current)
                           st <- get
                           trace ("Reboot on new examples: " ++ show newEx) $
                            stepOnAddedExprs $ map fst $ syntExprs st
                         else return maybeResult
                       else return Nothing

stepOnAddedExpr :: Expr -> SyntState (Maybe Expr)
stepOnAddedExpr expr = do rootGoal <- gets syntRoot
                          trySolveGoal expr rootGoal
                          maybeResult <- tryTerminate expr
                          if isJust maybeResult
                          then return maybeResult
                          else do
                            -- try resolve goals & resolvers, generate new resolvers
                            exprOutputs <- calcExprOutputs expr
                            -- NOTE: now done in fowardStep
                            -- when (foldl (compareExprOutputs exprOutputs) True $ map fst $ syntExprs st) $ modify $ \st -> st { syntExprs = tail $ syntExprs st }
                            unsolvedGoals <- gets syntUnsolvedGoals
                            foldM_ (const $ trySolveGoal expr) False unsolvedGoals -- solve existing goals
                            resolvers <- gets syntResolvers
                            foldM_ (const tryResolve) Nothing resolvers -- resolve existing goals
                            modify $ \st -> foldl (splitGoalsFold expr exprOutputs) st [syntRoot st] -- TODO: use Set.toList $ syntUnsolvedGoals st ?
                            tryTerminate expr
                            -- NOTE: replaced by tryTerminate
                            -- gets $ \st -> syntRoot st  `Map.lookup` syntSolvedGoals st
  where splitGoalsFold expr outputs st goal@(Goal expected) = let matches = zipWith matchResult outputs expected in
                                                              if not $ any (fromMaybe False) matches then st else
                                                              let matchesBool = map (fromMaybe True) matches in
                                                              execState (do r <- splitGoalStep goal matchesBool
                                                                            exprs <- gets syntExprs
                                                                            foldM_ (const $ flip trySolveGoal $ resolverCond r) False $ map fst exprs
                                                                            exprs <- gets syntExprs
                                                                            foldM_ (const $ flip trySolveGoal $ resolverElse r) False $ map fst exprs
                                                                            -- TODO: replace with always solve goal
                                                                            trySolveGoal expr (resolverThen r)) st
        matchResult :: Result Value -> Maybe Value -> Maybe Bool -- Nothing for unimportant matches marked as Nothing
        matchResult (NewExamples {}) _ = Just False
        matchResult _ Nothing = Nothing
        matchResult (RecError {}) _ = Just False
        matchResult (Result x) (Just y) = Just $ x == y
        -- compareExprOutputs outputs False _ = False
        -- compareExprOutputs outputs True e = do eOutputs <- calcExprOutputs e
        --                                        outputs == eOutputs

stepOnAddedExprs :: [Expr] -> SyntState (Maybe Expr)
stepOnAddedExprs = foldM step Nothing
  where step :: Maybe Expr -> Expr -> SyntState (Maybe Expr)
        step res@(Just {}) _ = return res
        step Nothing expr = stepOnAddedExpr expr

-- returns result and valid expr
stepOnNewExpr :: Expr -> [Expr] -> SyntState (Maybe Expr, Maybe Expr)
stepOnNewExpr comp args = do expr <- forwardStep comp args
                             case expr of
                               Just expr' -> do res <- stepOnAddedExpr expr'
                                                return (res, expr)
                               Nothing -> return (Nothing, Nothing)

-- stages:
-- init state
-- 1. gen new step exprs
-- 2. process exprs by one
-- 3. try terminate / saturate
-- 4. try to solve existing goals
-- 5. make resolutions if goals solved
-- 6. split goals, where expr partially matched
syntesisStep :: Int -> [[Expr]] -> SyntState (Maybe Expr)
syntesisStep 0 _ = return Nothing
syntesisStep steps prevExprs = -- oracle should be defined on the provided emample inputs
                               do let genExprs = genStep prevExprs
                                  (result, validExprs) <- foldM step (Nothing, []) genExprs
                                  if isJust result
                                  then return result
                                  else trace ("steps left: " ++ show (steps - 1)) $ syntesisStep (steps - 1) (validExprs : prevExprs)
  where step res@(Just {}, _) _ = return res
        step (Nothing, exprs) expr = do (res, val) <- uncurry stepOnNewExpr expr
                                        return (res, maybeToList val ++ exprs)

syntesis' :: [[Expr]] -> Int -> Oracle -> [[Value]] -> (Maybe Expr, Synt)
syntesis' exprs steps oracle inputs = -- oracle should be defined on the providid examples
                                      let outputs = map (fromMaybe undefined . oracleFunc oracle) inputs in
                                      runState (syntesisStep steps exprs) (createSynt oracle $ zip inputs outputs)

syntesis :: Int -> Oracle -> [[Value]] -> (Maybe Expr, Synt)
syntesis = syntesis' []

------ examples

mainExamples :: [[Value]]
mainExamples = [[ListV [IntV 1, IntV 2, IntV 3]]]

allExamples :: [[Value]]
-- allExamples = [[ListV [IntV 2, IntV 3]], [ListV [IntV 3]], [ListV []]]
allExamples = [[ListV [IntV 0, IntV 1, IntV 2, IntV 3]], [ListV [IntV 2, IntV 3]], [ListV [IntV 3]], [ListV []]]

listOracleOf :: OracleFunc -> Type -> Oracle
listOracleOf f t = Oracle { oracleTypes = TypeConf { typeConfInput = [ListT IntT],
                                                     typeConfOutput = t },
                            oracleFunc = f }

--- reverse

reverseOracle :: OracleFunc
-- reverseOracle [ListV xs] = Just $ ListV $ reverse xs
reverseOracle [ListV xs] | all isInt xs = Just $ ListV $ reverse xs
reverseOracle _ = Nothing

reverseExpr :: Expr
reverseExpr = IfE { ifCond = IsEmptyE (InputE 0),
                    ifDoThen = EmptyListE,
                    ifDoElse = SelfE [TailE (InputE 0)] :++: (HeadE (InputE 0) ::: EmptyListE) }

reverseConf :: Conf
reverseConf = Conf { confInput = head allExamples,
                     confOracle = listOracleOf reverseOracle $ ListT IntT,
                     confExamples = allExamples,
                     confTryFindExamples = True }

--- stutter

stutterOracle :: OracleFunc
stutterOracle [ListV (x : xs)] | isInt x = do ListV xs' <- stutterOracle [ListV xs]
                                              return $ ListV $ x : x : xs'
stutterOracle [ListV []] = Just $ ListV []
stutterOracle _ = Nothing

stutterExpr :: Expr
stutterExpr = IfE { ifCond = IsEmptyE (InputE 0),
                    ifDoThen = EmptyListE,
                    ifDoElse = HeadE (InputE 0) ::: (HeadE (InputE 0) ::: SelfE [TailE (InputE 0)]) }

stutterConf :: Conf
stutterConf = Conf { confInput = head allExamples,
                     confOracle = listOracleOf stutterOracle $ ListT IntT,
                     confExamples = allExamples,
                     confTryFindExamples = True }

--- length

lengthOracle :: OracleFunc
lengthOracle [ListV xs] = Just $ IntV $ length xs
lengthOracle _ = Nothing

lengthExpr :: Expr
lengthExpr = IfE { ifCond = IsEmptyE (InputE 0),
                    ifDoThen = ZeroE,
                    ifDoElse = IncE $ SelfE [TailE (InputE 0)] }

lengthConf :: Conf
lengthConf = Conf { confInput = head allExamples,
                     confOracle = listOracleOf lengthOracle IntT,
                     confExamples = allExamples,
                     confTryFindExamples = True }

---

idOracle :: OracleFunc
idOracle [x] = Just x
idOracle _ = Nothing

main = do steps <- readLn :: IO Int
          -- print $ fst $ syntesis steps (listOracleOf reverseOracle $ ListT IntT) allExamples
          print $ fst $ syntesis steps (listOracleOf stutterOracle $ ListT IntT) allExamples

-- main = print $ (SelfE (TailE (InputE ZeroE) ::: EmptyListE) :++: (HeadE (InputE ZeroE) ::: EmptyListE)) `elem` (map fst $ syntExprs $ snd $ syntesis 10 reverseOracle allExamples)
-- Just (IfE {ifCond = IsEmptyE (InputE ZeroE), ifDoThen = InputE ZeroE :++: TailE (InputE ZeroE :++: (InputE ZeroE :++: (ZeroE ::: EmptyListE))), ifDoElse = SelfE (TailE (InputE ZeroE) ::: EmptyListE) :++: (HeadE (InputE ZeroE) ::: EmptyListE)})