3 files changed, 148 insertions, 2 deletions

diff --git a/spec/inputs/in_expression.yue b/spec/inputs/in_expression.yue
index 6faff4e..7ee0816 100644
--- a/spec/inputs/in_expression.yue
+++ b/spec/inputs/in_expression.yue
@@ -35,5 +35,44 @@ do
         check -> return x in tb
         check -> x not in tb
 
+do
+        f1 = ->
+                a = 2
+                a in [
+                        1
+                        2
+                        3
+                        4
+                ]
+        f2 = ->
+                a not in [
+                        1
+                        2
+                        3
+                        4
+                ]
+        f3 = ->
+                a = 2
+                a in { b, c, d
+                        e, f
+                        g
+                }
+        f4 = ->
+                a = 2
+                b = a in [ 1, 2
+                        3, 4
+                ]
+        f5 = ->
+                a = 2
+                return a in [ 1, 2
+                        3, 4
+                ]
+        f6 = ->
+                b = a not in {
+                        1
+                        2
+                        x: 3
+                }
+
 nil
 
diff --git a/spec/outputs/in_expression.lua b/spec/outputs/in_expression.lua
index 04c08e7..afcc030 100644
--- a/spec/outputs/in_expression.lua
+++ b/spec/outputs/in_expression.lua
@@ -162,4 +162,53 @@ do
                 return true
         end)
 end
+do
+        local f1
+        f1 = function()
+                local a = 2
+                return (1 == a or 2 == a or 3 == a or 4 == a)
+        end
+        local f2
+        f2 = function()
+                local _val_0 = a
+                return not (1 == _val_0 or 2 == _val_0 or 3 == _val_0 or 4 == _val_0)
+        end
+        local f3
+        f3 = function()
+                local a = 2
+                return (b == a or c == a or d == a or e == a or f == a or g == a)
+        end
+        local f4
+        f4 = function()
+                local a = 2
+                local b
+                b = (1 == a or 2 == a or 3 == a or 4 == a)
+        end
+        local f5
+        f5 = function()
+                local a = 2
+                return (1 == a or 2 == a or 3 == a or 4 == a)
+        end
+        local f6
+        f6 = function()
+                local b
+                do
+                        local _check_0 = {
+                                1,
+                                2,
+                                x = 3
+                        }
+                        local _val_0 = a
+                        local _find_0 = false
+                        for _index_0 = 1, #_check_0 do
+                                local _item_0 = _check_0[_index_0]
+                                if _item_0 == _val_0 then
+                                        _find_0 = true
+                                        break
+                                end
+                        end
+                        b = not _find_0
+                end
+        end
+end
 return nil
diff --git a/src/yuescript/yue_compiler.cpp b/src/yuescript/yue_compiler.cpp
index ba74d91..3892d72 100644
--- a/src/yuescript/yue_compiler.cpp
+++ b/src/yuescript/yue_compiler.cpp
@@ -75,7 +75,7 @@ static std::unordered_set<std::string> Metamethods = {
         "close"s // Lua 5.4
 };
 
-const std::string_view version = "0.21.6"sv;
+const std::string_view version = "0.21.7"sv;
 const std::string_view extension = "yue"sv;
 
 class CompileError : public std::logic_error {
@@ -5946,7 +5946,52 @@ private:
                                         varName.clear();
                                 }
                         }
+                        ast_ptr<false, InDiscrete_t> discrete;
                         if (auto inExp = unary_exp->inExp->item.as<Exp_t>()) {
+                                BLOCK_START
+                                auto value = singleValueFrom(inExp);
+                                BREAK_IF(!value);
+                                auto sval = value->item.as<SimpleValue_t>();
+                                BREAK_IF(!sval);
+                                if (auto table = sval->value.as<TableLit_t>()) {
+                                        discrete = inExp->new_ptr<InDiscrete_t>();
+                                        for (ast_node* val : table->values.objects()) {
+                                                if (auto def = ast_cast<NormalDef_t>(val)) {
+                                                        if (def->defVal) {
+                                                                discrete = nullptr;
+                                                                break;
+                                                        } else {
+                                                                discrete->values.push_back(def->item);
+                                                        }
+                                                } else if (ast_is<Exp_t>(val)) {
+                                                        discrete->values.push_back(val);
+                                                } else {
+                                                        discrete = nullptr;
+                                                        break;
+                                                }
+                                        }
+                                } else if (auto comp = sval->value.as<Comprehension_t>()) {
+                                        if (comp->items.size() != 2 || !ast_is<CompInner_t>(comp->items.back())) {
+                                                discrete = inExp->new_ptr<InDiscrete_t>();
+                                                for (ast_node* val : comp->items.objects()) {
+                                                        if (auto def = ast_cast<NormalDef_t>(val)) {
+                                                                if (def->defVal) {
+                                                                        discrete = nullptr;
+                                                                        break;
+                                                                } else {
+                                                                        discrete->values.push_back(def->item);
+                                                                }
+                                                        } else {
+                                                                discrete = nullptr;
+                                                                break;
+                                                        }
+                                                }
+                                        }
+                                }
+                                BLOCK_END
+
+                                BLOCK_START
+                                BREAK_IF(discrete);
                                 str_list temp;
                                 auto checkVar = singleVariableFrom(inExp, false);
                                 if (usage == ExpUsage::Assignment) {
@@ -6069,8 +6114,11 @@ private:
                                         out.push_back(join(temp));
                                         return;
                                 }
+                                BLOCK_END
+                        }
+                        if (!discrete) {
+                                discrete = unary_exp->inExp->item.to<InDiscrete_t>();
                         }
-                        auto discrete = unary_exp->inExp->item.to<InDiscrete_t>();
                         if (usage == ExpUsage::Closure && discrete->values.size() == 1) {
                                 str_list tmp;
                                 transformExp(static_cast<Exp_t*>(discrete->values.front()), tmp, ExpUsage::Closure);
@@ -6151,6 +6199,13 @@ private:
                                 for (auto exp : discrete->values.objects()) {
                                         transformExp(static_cast<Exp_t*>(exp), tmp, ExpUsage::Closure);
                                 }
+                                if (usage == ExpUsage::Assignment) {
+                                        str_list tmpList;
+                                        transformExp(static_cast<Exp_t*>(assignList->exprs.front()), tmpList, ExpUsage::Closure);
+                                        _buf << indent() << tmpList.back() << " = "sv;
+                                } else if (usage == ExpUsage::Return) {
+                                        _buf << indent() << "return "sv;
+                                }
                                 if (unary_exp->inExp->not_) {
                                         _buf << "not "sv;
                                 }
@@ -6162,6 +6217,9 @@ private:
                                         }
                                 }
                                 _buf << ')';
+                                if (usage == ExpUsage::Assignment || usage == ExpUsage::Return) {
+                                        _buf << nll(discrete);
+                                }
                                 out.push_back(clearBuf());
                         }
                         return;