(with_import "./collections.kp"
(let (
    ; This partial eval only works via the root env,
    ; since we can't tell from outside if something takes an env
    ; or not
    ; Also, based on that, we have to switch on every primitive...


    ; Here is every form in k'
    ; True
    ; False
    ; Env: *KPEnv
    ; Combiner: KPCombiner / BuiltinCombiner: KPBuiltinCombiner
    ; String: str
    ; Symbol: str
    ; Int: int
    ; Array: rc<vec<KPValue>>
    ; Nil


    ; Ok, some more things we need / need to change

    ; meta...

    ; Env will need special checking for escaping, including when contained inside of an array or other combinator
    ; Actually, that extends to checking escaping for combinators in general, since they may close over laters

    ; Ok, instead of just ['now v] and ['later v], we have these marked values
    ; ['val v] (v can be an array)
    ; ['later c]
    ; ['marked_array a] (a contains marked values)
    ; ['comb wrap_level de? se params body <actual_function>]
    ; ['prim_comb <handler_function> <actual_function>]
    ; ['env [ ['symbol marked_value ]... ] <actual_env>]


    val? (lambda (x) (= 'val (idx x 0)))
    .val (lambda (x) (idx x 1))
    later? (lambda (x) (= 'later (idx x 0)))
    .later (lambda (x) (idx x 1))
    marked_array? (lambda (x) (= 'marked_array (idx x 0)))
    comb? (lambda (x) (= 'comb (idx x 0)))
    .comb (lambda (x) (slice x 1 -1))
    prim_comb? (lambda (x) (= 'prim_comb (idx x 0)))
    .prim_comb (lambda (x) (idx x 1))
    marked_env? (lambda (x) (= 'env (idx x 0)))
    .env_marked (lambda (x) (idx x 1))
    .env_real (lambda (x) (idx x 2))

    env-lookup-helper (rec-lambda recurse (dict key i fail success) (cond (and (= i (- (len dict) 1)) (= nil (idx dict i)))  (fail)
                                                                          (= i (- (len dict) 1))                             (recurse (idx (idx dict i) 1) key 0 fail success)
                                                                          (= key (idx (idx dict i) 0))                       (success (idx (idx dict i) 1))
                                                                          true                                               (recurse dict key (+ i 1) fail success)))
    env-lookup (lambda (env key) (env-lookup-helper (idx env 1) key 0 (lambda () (error (str key " not found in env " dict))) (lambda (x) x)))

    strip (rec-lambda recurse (x)
        (do (println "calling strip with " x)
        (cond (val? x)          (.val x)
              (later? x)        (.later x)
              (marked_array? x) (map recurse (idx x 1))
              (comb? x)         (idx x 6)
              (prim_comb? x)    (idx x 2)
              (makred_env? x)   (error "Env escaped to strip!")
        ))
    )

    ; GAH ok additionally
    ; partial_eval_helper will have to deal with combinator values (at least, primitives, I suspect all)
    ; as it might have to use them to reconstruct an expression on strip,
    ; and we will have to partially-eval previously strip'ped expressions when
    ; calling functions whose definition we partially-evaled, etc.
    ; partial_eval_helper always takes in unmarked expressions and makes marked ones, maintaining marked envs
    ;
    ; If indeed we have to keep track of non-primitive combinator values (which I think makes sense for stripping),
    ; we'll have to continually keep a map of values to their definition (we do this now!).

    partial_eval_helper (rec-lambda recurse (x env comb_to_mark_map)
        (cond   (= x true)                      [comb_to_mark_map ['val true ]]
                (= x false)                     [comb_to_mark_map ['val false]]
                (env? x)                        (error (str "called partial_eval with an env " x))
                (combiner? x)                   [comb_to_mark_map (get-value comb_to_mark_map x)]
                (string? x)                     [comb_to_mark_map ['val x]]
                (symbol? x)                     [comb_to_mark_map (env-lookup env x)]
                (int? x)                        [comb_to_mark_map ['val x]]
                (and (array? x) (= 0 (len x)))  (error "Partial eval on empty array")
                (array? x)                      (let ( [comb_to_mark_map comb] (recurse (idx x 0) env comb_to_mark_map) )
                                                                             ; it seems like even if it's later we should be able to eval some?
                                                                             ; Maybe there should be something between 'later and 'comb made in vau
                                                                             ; for those sorts of cases, but maybe it doesn't matter?
                                                     (cond (later? comb)     [comb_to_mark_map ['later x]]
                                                           (prim_comb? comb) ((.prim_comb comb) env comb_to_mark_map (slice x 1 -1))
                                                           (comb? comb)      (let (
                                                                                 [wrap_level de? se params body actual_function] (.comb comb)
                                                                                 [comb_to_mark_map appropriatly_evaled_params] ((rec-lambda param-recurse (wrap params comb_to_mark_map)
                                                                                        (if (!= 0 wrap)
                                                                                            (let ([comb_to_mark_map evaled_params]
                                                                                                                    (foldl (lambda ([comb_to_mark_map ac] p)
                                                                                                                        (let ([comb_to_mark_map p] (recurse p env comb_to_mark_map))
                                                                                                                              [comb_to_mark_map (concat ac [p])]))
                                                                                                                    [comb_to_mark_map []]
                                                                                                                    (map strip params)))
                                                                                                 (param-recurse (- wrap 1) evaled_params comb_to_mark_map))
                                                                                            [comb_to_mark_map params])
                                                                                            ) wrap_level (map (lambda (p) ['val p]) (slice x 1 -1)) comb_to_mark_map)
                                                                                 de_entry (if (!= nil de?) [ [de? env] ]
                                                                                                           [])
                                                                                 _ (println "appropriately evaled params " appropriatly_evaled_params)
                                                                                 inner_env ['env (concat (zip params appropriatly_evaled_params) de_entry [se]) nil]
                                                                                 ;_ (println "inner_env is " inner_env)
                                                                             ) (recurse body inner_env comb_to_mark_map))
                                                           true              (error (str "Partial eval noticed that you will likely call not a function " x))))
                (nil? x)                        [comb_to_mark_map ['val x]]
                true                            (error (str "impossible partial_eval value " x))
        )
    )
    needs_params_val_lambda (vau de (f_sym) (let (
        actual_function (eval f_sym de)
        handler (lambda (de comb_to_mark_map params) (let (
                [comb_to_mark_map evaled_params] (foldl (lambda ([comb_to_mark_map evaleds] x) (let (
                                                                     [comb_to_mark_map evaled] (partial_eval_helper x de comb_to_mark_map)
                                                                 ) [comb_to_mark_map (cons evaled evaleds)])) [comb_to_mark_map []] params)
            )
            (if (foldl (lambda (a x) (and a (val? x))) true evaled_params) [comb_to_mark_map ['val (lapply actual_function (map .val evaled_params))]]
                                                                           [comb_to_mark_map ['later (cons actual_function (map strip evaled_params))]])))
        ) [f_sym ['prim_comb handler actual_function]]))
    give_up (vau de (f_sym) (let (
        actual_function (eval f_sym de)
        handler (lambda (de comb_to_mark_map params) [comb_to_mark_map ['later (cons actual_function params)]])
        ) [f_sym ['prim_comb handler actual_function]]))


    ; Our job is made a lot easier by the fact that these will
    ; all be stripped and we only care about symbols and things that could contain symbols,
    ; namely arrays
    ;
    ; This is a conservative analysis, since we can't always tell what constructs introduce
    ; a new binding scope & would be shadowing... we should at least be able to implement it for
    ; vau/lambda, but we won't at first
    closes_over_outside_vars (rec-lambda recurse (env body) (cond
                (symbol? body)                     (env-lookup-helper (idx env 1) body 0 (lambda () false) (lambda (x) true))
                (array? body)                      (foldl (lambda (a x) (or a (recurse env x))) false body)
                true                               false))

    root_marked_env ['env [
        ; Ok, so for combinators, it should partial eval the body.
        ; It should then check to see if the partial-evaled body has closed over
        ; any 'later values from above the combinator. If so, the combinator should
        ; evaluate to a ['later [vau de? params (strip partially_evaled_body)]], otherwise it can evaluate to a 'comb.
        ; Note that this 'later may be re-evaluated later if the parent function is called.
        ;(give_up vau)
        ['vau ['prim_comb (lambda (de comb_to_mark_map params) (let (
            de?         (if (= 3 (len params)) (idx params 0))
            vau_de?     (if (= nil de?) [] [de?])
            vau_params  (if (= nil de?) (idx params 0) (idx params 1))
            body        (if (= nil de?) (idx params 1) (idx params 2))
            inner_env   ['env (concat (map (lambda (p) [p ['later p]]) vau_params) (if (= nil de?) [] [ [de? ['later de?]] ]) [de]) nil]
            [comb_to_mark_map pe_body] (partial_eval_helper body inner_env comb_to_mark_map)
            spe_body (strip pe_body)
        ) (if (or (= nil (.env_real de)) (closes_over_outside_vars de spe_body)) [comb_to_mark_map ['later (concat [vau] vau_de? [vau_params spe_body])]]
                                                                                 (let (real_func (eval (concat [vau] vau_de? [vau_params spe_body]) (.env_real de))
                                                                                       marked_func ['comb 0 de? de vau_params spe_body real_func]
                                                                                 ) [(put comb_to_mark_map real_func marked_func) marked_func])))
        ) vau]]
        ['wrap ['prim_comb (lambda (de comb_to_mark_map params) (let (
                _ (if (!= 1 (len params)) (error (str "bad number of params to partial-eval wrap " params)))
                [comb_to_mark_map evaled] (partial_eval_helper (idx params 0) de comb_to_mark_map)
                ;_ (println "wrap evaled is " evaled)
            ) (if (comb? evaled) (let ([wrap_level de? se params body actual_function] (.comb evaled)
                                        wrapped_actual_fun (wrap actual_function)
                                        wrapped_marked_fun ['comb (+ 1 wrap_level) de? se params body wrapped_actual_fun]
                                      ) [(put comb_to_mark_map wrapped_actual_fun wrapped_marked_fun) wrapped_marked_fun])
                                 [comb_to_mark_map ['later [wrap (strip evaled)]]]))
        ) wrap]]
        ['unwrap ['prim_comb (lambda (de comb_to_mark_map params) (let (
                _ (if (!= 1 (len params)) (error (str "bad number of params to partial-eval unwrap " params)))
                [comb_to_mark_map evaled] (partial_eval_helper (idx params 0) de comb_to_mark_map)
                ;_ (println "unwrap evaled is " evaled)
            ) (if (comb? evaled) (let ([wrap_level de? se params body actual_function] (.comb evaled)
                                        unwrapped_actual_fun (unwrap actual_function)
                                        unwrapped_marked_fun ['comb (- wrap_level 1) de? se params body unwrapped_actual_fun]
                                      ) [(put comb_to_mark_map unwrapped_actual_fun unwrapped_marked_fun) unwrapped_marked_fun])
                                 [comb_to_mark_map ['later [unwrap (strip evaled)]]]))
        ) unwrap]]

        ; eval should have it's parameters partially -evaled, then partially-eval e again.
        ; failure can 'later at either point
        (give_up eval)
        (give_up cond)
        (needs_params_val_lambda symbol?)
        (needs_params_val_lambda int?)
        (needs_params_val_lambda string?)
        (give_up combiner?)
        (give_up env?)
        (needs_params_val_lambda nil?)
        (needs_params_val_lambda bool?)
        (give_up array?)
        (needs_params_val_lambda str-to-symbol)
        (needs_params_val_lambda get-text)
        (give_up array)
        (give_up len)
        (give_up idx)
        (give_up slice)
        (give_up concat)
        (needs_params_val_lambda +)
        (needs_params_val_lambda -)
        (needs_params_val_lambda *)
        (needs_params_val_lambda /)
        (needs_params_val_lambda %)
        (needs_params_val_lambda &)
        (needs_params_val_lambda |)
        (needs_params_val_lambda <<)
        (needs_params_val_lambda >>)
        (needs_params_val_lambda =)
        (needs_params_val_lambda !=)
        (needs_params_val_lambda <)
        (needs_params_val_lambda <=)
        (needs_params_val_lambda >)
        (needs_params_val_lambda >=)

        ; Don't forget, these short-circut with the truthy/falsey values
        (give_up and)
        (give_up or)

        ; pr-str
        (needs_params_val_lambda str)
        (needs_params_val_lambda prn)
        (give_up println)
        (give_up meta)
        (give_up with-meta)
        (give_up error)
        (give_up recover)
        (give_up read-string)
        (give_up slurp)
        (give_up get_line)
        (give_up write_file)
        ['empty_env ['env [] empty_env]]
        nil
    ] root_env]

    comb_to_mark_map (foldl (lambda (a x) (cond (= nil x)              a
                                                (comb?      (idx x 1)) (put a (idx (idx x 1) 6) (idx x 1))
                                                (prim_comb? (idx x 1)) (put a (idx (idx x 1) 2) (idx x 1))
                                                true                   a
                                          ) ) empty_dict (idx root_marked_env 1))
    partial_eval (lambda (x) (partial_eval_helper x root_marked_env comb_to_mark_map))
)
    (provide partial_eval strip)
))