TransPyC/lib/vast_c/vast_c.c

/* vast_c - C 版 Viper/Python 语法解析器（vast 加速层）
 *
 * 阶段 0：最小可行性验证——Tokenizer + 表达式/赋值/模块解析，
 * 直接构造 lib.vast.Nodes 的节点类实例，输出与 vast.parse 结构完全一致。
 *
 * 编译（MSYS2 MinGW gcc，Python 3.12，注意 -l 必须在源文件之后）：
 *   gcc -shared -O2 -std=c11 -I D:\Python312\include lib\vast_c\vast_c.c \
 *       -L D:\Python312\libs -lpython312 -o lib\vast_c\_parser.pyd
 *
 * 使用：
 *   from lib.vast_c import parse   # 由 lib/vast_c/__init__.py 导出
 *   tree = parse("x = 1")          # 与 vast.parse 同结构
 */
#define PY_SSIZE_T_CLEAN
#include <Python.h>

#include <stdbool.h>
#include <stdlib.h>
#include <string.h>

/* ================= 记号（Token） ================= */
typedef enum {
    TK_EOF = 0,
    TK_NAME,
    TK_NUMBER,
    TK_EQUAL,
    TK_LPAREN,
    TK_RPAREN,
    TK_NEWLINE,
} TokenKind;

typedef struct {
    TokenKind Kind;
    const char *Start;
    Py_ssize_t Len;
    int Line;
    int Col;
} Token;

/* ================= 解析器（Parser） ================= */
typedef struct {
    const char *Src;
    Py_ssize_t Len;
    Py_ssize_t Pos;
    int Line;
    int Col;
    Token Cur;
    int Error;
} Parser;

/* vast.nodes 模块与节点类缓存（进程级惰性初始化） */
static PyObject *s_Nodes = NULL;
static PyObject *s_Classes = NULL; /* dict: 类名 -> 类对象 */

static PyObject *NodesModule(void) {
    if (s_Nodes == NULL) {
        s_Nodes = PyImport_ImportModule("lib.vast.Nodes");
    }
    return s_Nodes;
}

/* 获取节点类对象（borrowed 引用，内部缓存避免重复属性查找） */
static PyObject *ClassGet(Parser *P, const char *Name) {
    PyObject *Mod = NodesModule();
    if (Mod == NULL) {
        return NULL;
    }
    if (s_Classes == NULL) {
        s_Classes = PyDict_New();
        if (s_Classes == NULL) {
            return NULL;
        }
    }
    PyObject *Key = PyUnicode_FromString(Name);
    if (Key == NULL) {
        return NULL;
    }
    PyObject *Cls = PyDict_GetItemWithError(s_Classes, Key); /* borrowed */
    if (Cls == NULL && !PyErr_Occurred()) {
        Cls = PyObject_GetAttrString(Mod, Name);
        if (Cls == NULL) {
            Py_DECREF(Key);
            return NULL;
        }
        if (PyDict_SetItem(s_Classes, Key, Cls) < 0) {
            Py_DECREF(Cls);
            Py_DECREF(Key);
            return NULL;
        }
    }
    Py_DECREF(Key);
    return Cls; /* borrowed */
}

/* 构造节点：按位置参数数组调用类（vast.nodes 的 __init__(*args, **kwargs)） */
static PyObject *NodeMake(Parser *P, const char *Name, PyObject **Args, Py_ssize_t NArgs) {
    PyObject *Cls = ClassGet(P, Name);
    if (Cls == NULL) {
        return NULL;
    }
    return PyObject_Vectorcall(Cls, Args, NArgs, NULL);
}

/* 设置位置属性（vast.parse 默认携带 lineno/col_offset/end_lineno/end_col_offset） */
static void NodePos(Parser *P, PyObject *Node, int Line, int Col) {
    PyObject *LineObj = PyLong_FromLong(Line);
    PyObject *ColObj = PyLong_FromLong(Col);
    PyObject *EndLineObj = PyLong_FromLong(P->Line);
    PyObject *EndColObj = PyLong_FromLong(P->Col);
    if (LineObj && ColObj && EndLineObj && EndColObj) {
        PyObject_SetAttrString(Node, "lineno", LineObj);
        PyObject_SetAttrString(Node, "col_offset", ColObj);
        PyObject_SetAttrString(Node, "end_lineno", EndLineObj);
        PyObject_SetAttrString(Node, "end_col_offset", EndColObj);
    }
    Py_XDECREF(LineObj);
    Py_XDECREF(ColObj);
    Py_XDECREF(EndLineObj);
    Py_XDECREF(EndColObj);
}

/* ================= Tokenizer ================= */

static int IsNameStart(int C) {
    return (C >= 'a' && C <= 'z') || (C >= 'A' && C <= 'Z') || C == '_';
}

static int IsNameChar(int C) {
    return IsNameStart(C) || (C >= '0' && C <= '9');
}

static int IsDigit(int C) {
    return C >= '0' && C <= '9';
}

static void NextToken(Parser *P) {
    /* 跳过空白（不含换行） */
    while (P->Pos < P->Len) {
        char C = P->Src[P->Pos];
        if (C == ' ' || C == '\t' || C == '\r') {
            P->Pos++;
            P->Col++;
        } else {
            break;
        }
    }
    Token T;
    T.Start = P->Src + P->Pos;
    T.Len = 0;
    T.Line = P->Line;
    T.Col = P->Col;
    if (P->Pos >= P->Len) {
        T.Kind = TK_EOF;
        P->Cur = T;
        return;
    }
    char C = P->Src[P->Pos];
    if (C == '\n') {
        T.Kind = TK_NEWLINE;
        T.Len = 1;
        P->Pos++;
        P->Line++;
        P->Col = 0;
    } else if (IsNameStart(C)) {
        Py_ssize_t Start = P->Pos;
        while (P->Pos < P->Len && IsNameChar(P->Src[P->Pos])) {
            P->Pos++;
        }
        T.Kind = TK_NAME;
        T.Len = P->Pos - Start;
        P->Col += (int)T.Len;
    } else if (IsDigit(C)) {
        Py_ssize_t Start = P->Pos;
        while (P->Pos < P->Len && (IsDigit(P->Src[P->Pos]) || P->Src[P->Pos] == '.')) {
            P->Pos++;
        }
        T.Kind = TK_NUMBER;
        T.Len = P->Pos - Start;
        P->Col += (int)T.Len;
    } else {
        switch (C) {
        case '=':
            T.Kind = TK_EQUAL;
            T.Len = 1;
            P->Pos++;
            P->Col++;
            break;
        case '(':
            T.Kind = TK_LPAREN;
            T.Len = 1;
            P->Pos++;
            P->Col++;
            break;
        case ')':
            T.Kind = TK_RPAREN;
            T.Len = 1;
            P->Pos++;
            P->Col++;
            break;
        default:
            T.Kind = TK_EOF;
            P->Error = 1;
            break;
        }
    }
    P->Cur = T;
}

/* ================= Parser ================= */

static PyObject *ParseExpression(Parser *P);

/* 原子表达式：数字常量 / 名字 */
static PyObject *ParsePrimary(Parser *P) {
    if (P->Cur.Kind == TK_NUMBER) {
        const char *Text = P->Cur.Start;
        Py_ssize_t Len = P->Cur.Len;
        int Line = P->Cur.Line;
        int Col = P->Cur.Col;
        int IsFloat = 0;
        for (Py_ssize_t I = 0; I < Len; I++) {
            if (Text[I] == '.') {
                IsFloat = 1;
                break;
            }
        }
        PyObject *StrObj = PyUnicode_FromStringAndSize(Text, Len);
        if (StrObj == NULL) {
            P->Error = 1;
            return NULL;
        }
        PyObject *ValObj;
        if (IsFloat) {
            ValObj = PyFloat_FromString(StrObj);
        } else {
            ValObj = PyLong_FromUnicodeObject(StrObj, 10);
        }
        Py_DECREF(StrObj);
        if (ValObj == NULL) {
            P->Error = 1;
            return NULL;
        }
        PyObject *Args[2];
        Args[0] = ValObj;
        Args[1] = Py_None;
        PyObject *Node = NodeMake(P, "Constant", Args, 2);
        Py_DECREF(ValObj);
        if (Node == NULL) {
            return NULL;
        }
        NodePos(P, Node, Line, Col);
        NextToken(P);
        return Node;
    }
    if (P->Cur.Kind == TK_NAME) {
        const char *Text = P->Cur.Start;
        Py_ssize_t Len = P->Cur.Len;
        int Line = P->Cur.Line;
        int Col = P->Cur.Col;
        PyObject *LoadCls = ClassGet(P, "Load");
        if (LoadCls == NULL) {
            return NULL;
        }
        PyObject *Ctx = PyObject_Vectorcall(LoadCls, NULL, 0, NULL);
        PyObject *Id = PyUnicode_FromStringAndSize(Text, Len);
        if (Ctx == NULL || Id == NULL) {
            Py_XDECREF(Ctx);
            Py_XDECREF(Id);
            return NULL;
        }
        PyObject *Args[2];
        Args[0] = Id;
        Args[1] = Ctx;
        PyObject *Node = NodeMake(P, "Name", Args, 2);
        Py_DECREF(Ctx);
        Py_DECREF(Id);
        if (Node == NULL) {
            return NULL;
        }
        NodePos(P, Node, Line, Col);
        NextToken(P);
        return Node;
    }
    P->Error = 1;
    return NULL;
}

/* 表达式（阶段 0：仅原子表达式） */
static PyObject *ParseExpression(Parser *P) {
    return ParsePrimary(P);
}

/* 语句：Name '=' expr → Assign；其余表达式 → Expr 语句 */
static PyObject *ParseStatement(Parser *P) {
    int Line = P->Cur.Line;
    int Col = P->Cur.Col;
    if (P->Cur.Kind == TK_NAME) {
        /* 向前看是否为赋值：Name '=' ... */
        const char *Text = P->Cur.Start;
        Py_ssize_t Len = P->Cur.Len;
        NextToken(P);
        if (P->Cur.Kind == TK_EQUAL) {
            NextToken(P);
            PyObject *Value = ParseExpression(P);
            if (Value == NULL) {
                return NULL;
            }
            PyObject *StoreCls = ClassGet(P, "Store");
            if (StoreCls == NULL) {
                Py_DECREF(Value);
                return NULL;
            }
            PyObject *Ctx = PyObject_Vectorcall(StoreCls, NULL, 0, NULL);
            PyObject *Id = PyUnicode_FromStringAndSize(Text, Len);
            if (Ctx == NULL || Id == NULL) {
                Py_XDECREF(Ctx);
                Py_XDECREF(Id);
                Py_DECREF(Value);
                return NULL;
            }
            PyObject *TargetArgs[2];
            TargetArgs[0] = Id;
            TargetArgs[1] = Ctx;
            PyObject *Target = NodeMake(P, "Name", TargetArgs, 2);
            Py_DECREF(Ctx);
            Py_DECREF(Id);
            if (Target == NULL) {
                Py_DECREF(Value);
                return NULL;
            }
            NodePos(P, Target, Line, Col);
            PyObject *Targets = PyList_New(1);
            PyList_SetItem(Targets, 0, Target); /* 窃取引用 */
            PyObject *AssignArgs[3];
            AssignArgs[0] = Targets;
            AssignArgs[1] = Value;
            AssignArgs[2] = Py_None; /* type_comment */
            PyObject *Node = NodeMake(P, "Assign", AssignArgs, 3);
            Py_DECREF(Value);
            if (Node == NULL) {
                Py_DECREF(Targets);
                return NULL;
            }
            NodePos(P, Node, Line, Col);
            return Node;
        }
        /* 非赋值：Cur 已前进，直接用保存的文本构造 Load Name 表达式语句 */
        PyObject *LoadCls = ClassGet(P, "Load");
        if (LoadCls == NULL) {
            return NULL;
        }
        PyObject *Ctx = PyObject_Vectorcall(LoadCls, NULL, 0, NULL);
        PyObject *Id = PyUnicode_FromStringAndSize(Text, Len);
        if (Ctx == NULL || Id == NULL) {
            Py_XDECREF(Ctx);
            Py_XDECREF(Id);
            return NULL;
        }
        PyObject *NameArgs[2];
        NameArgs[0] = Id;
        NameArgs[1] = Ctx;
        PyObject *NameNode = NodeMake(P, "Name", NameArgs, 2);
        Py_DECREF(Ctx);
        Py_DECREF(Id);
        if (NameNode == NULL) {
            return NULL;
        }
        NodePos(P, NameNode, Line, Col);
        PyObject *ExprArgs[1];
        ExprArgs[0] = NameNode;
        PyObject *Node = NodeMake(P, "Expr", ExprArgs, 1);
        Py_DECREF(NameNode);
        if (Node == NULL) {
            return NULL;
        }
        NodePos(P, Node, Line, Col);
        return Node;
    }
    /* 一般表达式语句 */
    PyObject *ExprVal = ParsePrimary(P);
    if (ExprVal == NULL) {
        return NULL;
    }
    PyObject *ExprArgs[1];
    ExprArgs[0] = ExprVal;
    PyObject *Node = NodeMake(P, "Expr", ExprArgs, 1);
    Py_DECREF(ExprVal);
    if (Node == NULL) {
        return NULL;
    }
    NodePos(P, Node, Line, Col);
    return Node;
}

/* 模块：语句序列 */
static PyObject *ParseModule(Parser *P) {
    PyObject *Body = PyList_New(0);
    if (Body == NULL) {
        return NULL;
    }
    NextToken(P);
    while (P->Cur.Kind != TK_EOF) {
        if (P->Cur.Kind == TK_NEWLINE) {
            NextToken(P);
            continue;
        }
        PyObject *Stmt = ParseStatement(P);
        if (Stmt == NULL) {
            break;
        }
        PyList_Append(Body, Stmt);
        Py_DECREF(Stmt);
    }
    PyObject *Ti = PyList_New(0);
    if (Ti == NULL) {
        Py_DECREF(Body);
        return NULL;
    }
    PyObject *Args[2];
    Args[0] = Body;
    Args[1] = Ti;
    PyObject *Mod = NodeMake(P, "Module", Args, 2);
    return Mod;
}

/* ================= 模块入口 ================= */

static PyObject *VastC_Parse(PyObject *Self, PyObject *Args, PyObject *Kwargs) {
    static char *KwList[] = {"source", "mode", NULL};
    const char *Src = NULL;
    Py_ssize_t Len = 0;
    const char *Mode = "exec";
    if (!PyArg_ParseTupleAndKeywords(Args, Kwargs, "s#|s:parse", KwList, &Src, &Len, &Mode)) {
        return NULL;
    }
    Parser P;
    P.Src = Src;
    P.Len = Len;
    P.Pos = 0;
    P.Line = 1;
    P.Col = 0;
    P.Error = 0;
    PyObject *Tree = ParseModule(&P);
    if (Tree == NULL) {
        if (P.Error) {
            PyErr_SetString(PyExc_SyntaxError, "vast_c: 语法错误（阶段 0 仅支持赋值/名字/数字）");
        }
        return NULL;
    }
    if (strcmp(Mode, "exec") == 0) {
        return Tree;
    }
    if (strcmp(Mode, "eval") == 0) {
        /* 复用 vast.parse 的 eval 语义：Expose Expr 的 value 为 Expression 节点 */
        PyObject *BodyAttr = PyObject_GetAttrString(Tree, "body");
        Py_DECREF(Tree);
        if (BodyAttr == NULL || PyList_GET_SIZE(BodyAttr) != 1) {
            Py_XDECREF(BodyAttr);
            PyErr_SetString(PyExc_SyntaxError, "vast_c: eval 模式需要单一表达式");
            return NULL;
        }
        PyObject *Stmt = PyList_GetItem(BodyAttr, 0); /* borrowed */
        PyObject *Value = PyObject_GetAttrString(Stmt, "value");
        Py_DECREF(BodyAttr);
        if (Value == NULL) {
            return NULL;
        }
        PyObject *Args2[1];
        Args2[0] = Value;
        PyObject *ExprNode = NodeMake(&P, "Expression", Args2, 1);
        Py_DECREF(Value);
        return ExprNode;
    }
    Py_DECREF(Tree);
    PyErr_SetString(PyExc_ValueError, "vast_c: 不支持的 mode");
    return NULL;
}

static PyMethodDef VastC_Methods[] = {
    {"parse", (PyCFunction)VastC_Parse, METH_VARARGS | METH_KEYWORDS,
     "parse(source, mode='exec') -> vast.nodes 树（阶段 0 子集）"},
    {NULL, NULL, 0, NULL},
};

static struct PyModuleDef VastC_Module = {
    PyModuleDef_HEAD_INIT,
    "vast_c",
    "vast_c - C 版 Viper/Python 解析器（vast 加速层）",
    -1,
    VastC_Methods,
};

__declspec(dllexport) PyMODINIT_FUNC PyInit__parser(void) {
    return PyModule_Create(&VastC_Module);
}